Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -26,7 +26,6 @@ from html_templates import (
|
|
26 |
format_description_html,
|
27 |
format_single_dog_result,
|
28 |
format_multiple_breeds_result,
|
29 |
-
format_error_message,
|
30 |
format_unknown_breed_message,
|
31 |
format_not_dog_message,
|
32 |
format_warning_html,
|
@@ -240,37 +239,6 @@ def predict_single_dog(image):
|
|
240 |
|
241 |
return probabilities[0], breeds[:3], relative_probs
|
242 |
|
243 |
-
# @spaces.GPU
|
244 |
-
# def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
|
245 |
-
|
246 |
-
# results = model_manager.yolo_model(image, conf=conf_threshold,
|
247 |
-
# iou=iou_threshold)[0]
|
248 |
-
|
249 |
-
# dogs = []
|
250 |
-
# boxes = []
|
251 |
-
# for box in results.boxes:
|
252 |
-
# if box.cls == 16: # COCO dataset class for dog is 16
|
253 |
-
# xyxy = box.xyxy[0].tolist()
|
254 |
-
# confidence = box.conf.item()
|
255 |
-
# boxes.append((xyxy, confidence))
|
256 |
-
|
257 |
-
# if not boxes:
|
258 |
-
# dogs.append((image, 1.0, [0, 0, image.width, image.height]))
|
259 |
-
# else:
|
260 |
-
# nms_boxes = non_max_suppression(boxes, iou_threshold)
|
261 |
-
|
262 |
-
# for box, confidence in nms_boxes:
|
263 |
-
# x1, y1, x2, y2 = box
|
264 |
-
# w, h = x2 - x1, y2 - y1
|
265 |
-
# x1 = max(0, x1 - w * 0.05)
|
266 |
-
# y1 = max(0, y1 - h * 0.05)
|
267 |
-
# x2 = min(image.width, x2 + w * 0.05)
|
268 |
-
# y2 = min(image.height, y2 + h * 0.05)
|
269 |
-
# cropped_image = image.crop((x1, y1, x2, y2))
|
270 |
-
# dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
|
271 |
-
|
272 |
-
# return dogs
|
273 |
-
|
274 |
@spaces.GPU
|
275 |
def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
|
276 |
"""
|
@@ -374,131 +342,13 @@ def create_breed_comparison(breed1: str, breed2: str) -> dict:
|
|
374 |
|
375 |
return comparison_data
|
376 |
|
377 |
-
|
378 |
-
# def predict(image):
|
379 |
-
# """
|
380 |
-
# Main prediction function that handles both single and multiple dog detection.
|
381 |
-
|
382 |
-
# Args:
|
383 |
-
# image: PIL Image or numpy array
|
384 |
-
|
385 |
-
# Returns:
|
386 |
-
# tuple: (html_output, annotated_image, initial_state)
|
387 |
-
# """
|
388 |
-
|
389 |
-
# if image is None:
|
390 |
-
# return format_warning_html("Please upload an image to start."), None, None
|
391 |
-
|
392 |
-
# try:
|
393 |
-
# if isinstance(image, np.ndarray):
|
394 |
-
# image = Image.fromarray(image)
|
395 |
-
|
396 |
-
# # Detect dogs in the image
|
397 |
-
# dogs = detect_multiple_dogs(image)
|
398 |
-
# color_scheme = get_color_scheme(len(dogs) == 1)
|
399 |
-
|
400 |
-
# # Prepare for annotation
|
401 |
-
# annotated_image = image.copy()
|
402 |
-
# draw = ImageDraw.Draw(annotated_image)
|
403 |
-
|
404 |
-
# try:
|
405 |
-
# font = ImageFont.truetype("arial.ttf", 24)
|
406 |
-
# except:
|
407 |
-
# font = ImageFont.load_default()
|
408 |
-
|
409 |
-
# dogs_info = ""
|
410 |
-
|
411 |
-
# # Process each detected dog
|
412 |
-
# for i, (cropped_image, detection_confidence, box) in enumerate(dogs):
|
413 |
-
# color = color_scheme if len(dogs) == 1 else color_scheme[i % len(color_scheme)]
|
414 |
-
|
415 |
-
# # Draw box and label on image
|
416 |
-
# draw.rectangle(box, outline=color, width=4)
|
417 |
-
# label = f"Dog {i+1}"
|
418 |
-
# label_bbox = draw.textbbox((0, 0), label, font=font)
|
419 |
-
# label_width = label_bbox[2] - label_bbox[0]
|
420 |
-
# label_height = label_bbox[3] - label_bbox[1]
|
421 |
-
|
422 |
-
# # Draw label background and text
|
423 |
-
# label_x = box[0] + 5
|
424 |
-
# label_y = box[1] + 5
|
425 |
-
# draw.rectangle(
|
426 |
-
# [label_x - 2, label_y - 2, label_x + label_width + 4, label_y + label_height + 4],
|
427 |
-
# fill='white',
|
428 |
-
# outline=color,
|
429 |
-
# width=2
|
430 |
-
# )
|
431 |
-
# draw.text((label_x, label_y), label, fill=color, font=font)
|
432 |
-
|
433 |
-
# # Predict breed
|
434 |
-
# top1_prob, topk_breeds, relative_probs = predict_single_dog(cropped_image)
|
435 |
-
# combined_confidence = detection_confidence * top1_prob
|
436 |
-
|
437 |
-
# # Format results based on confidence with error handling
|
438 |
-
# try:
|
439 |
-
# if combined_confidence < 0.2:
|
440 |
-
# dogs_info += format_error_message(color, i+1)
|
441 |
-
# elif top1_prob >= 0.45:
|
442 |
-
# breed = topk_breeds[0]
|
443 |
-
# description = get_dog_description(breed)
|
444 |
-
# # Handle missing breed description
|
445 |
-
# if description is None:
|
446 |
-
# # 如果沒有描述,創建一個基本描述
|
447 |
-
# description = {
|
448 |
-
# "Name": breed,
|
449 |
-
# "Size": "Unknown",
|
450 |
-
# "Exercise Needs": "Unknown",
|
451 |
-
# "Grooming Needs": "Unknown",
|
452 |
-
# "Care Level": "Unknown",
|
453 |
-
# "Good with Children": "Unknown",
|
454 |
-
# "Description": f"Identified as {breed.replace('_', ' ')}"
|
455 |
-
# }
|
456 |
-
# dogs_info += format_single_dog_result(breed, description, color)
|
457 |
-
# else:
|
458 |
-
# # 修改format_multiple_breeds_result的調用,包含錯誤���理
|
459 |
-
# dogs_info += format_multiple_breeds_result(
|
460 |
-
# topk_breeds,
|
461 |
-
# relative_probs,
|
462 |
-
# color,
|
463 |
-
# i+1,
|
464 |
-
# lambda breed: get_dog_description(breed) or {
|
465 |
-
# "Name": breed,
|
466 |
-
# "Size": "Unknown",
|
467 |
-
# "Exercise Needs": "Unknown",
|
468 |
-
# "Grooming Needs": "Unknown",
|
469 |
-
# "Care Level": "Unknown",
|
470 |
-
# "Good with Children": "Unknown",
|
471 |
-
# "Description": f"Identified as {breed.replace('_', ' ')}"
|
472 |
-
# }
|
473 |
-
# )
|
474 |
-
# except Exception as e:
|
475 |
-
# print(f"Error formatting results for dog {i+1}: {str(e)}")
|
476 |
-
# dogs_info += format_error_message(color, i+1)
|
477 |
-
|
478 |
-
# # Wrap final HTML output
|
479 |
-
# html_output = format_multi_dog_container(dogs_info)
|
480 |
-
|
481 |
-
# # Prepare initial state
|
482 |
-
# initial_state = {
|
483 |
-
# "dogs_info": dogs_info,
|
484 |
-
# "image": annotated_image,
|
485 |
-
# "is_multi_dog": len(dogs) > 1,
|
486 |
-
# "html_output": html_output
|
487 |
-
# }
|
488 |
-
|
489 |
-
# return html_output, annotated_image, initial_state
|
490 |
-
|
491 |
-
# except Exception as e:
|
492 |
-
# error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
493 |
-
# print(error_msg)
|
494 |
-
# return format_warning_html(error_msg), None, None
|
495 |
-
|
496 |
|
497 |
@spaces.GPU
|
498 |
def predict(image):
|
499 |
"""
|
500 |
主要的預測函數,負責處理狗的檢測和品種辨識。
|
501 |
它整合了YOLO的物體檢測和專門的品種分類模型。
|
|
|
502 |
|
503 |
Args:
|
504 |
image: PIL Image 或 numpy array
|
|
|
26 |
format_description_html,
|
27 |
format_single_dog_result,
|
28 |
format_multiple_breeds_result,
|
|
|
29 |
format_unknown_breed_message,
|
30 |
format_not_dog_message,
|
31 |
format_warning_html,
|
|
|
239 |
|
240 |
return probabilities[0], breeds[:3], relative_probs
|
241 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
242 |
@spaces.GPU
|
243 |
def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
|
244 |
"""
|
|
|
342 |
|
343 |
return comparison_data
|
344 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
345 |
|
346 |
@spaces.GPU
|
347 |
def predict(image):
|
348 |
"""
|
349 |
主要的預測函數,負責處理狗的檢測和品種辨識。
|
350 |
它整合了YOLO的物體檢測和專門的品種分類模型。
|
351 |
+
實施雙層檢測,非狗會直接忽略.
|
352 |
|
353 |
Args:
|
354 |
image: PIL Image 或 numpy array
|