Update app.py

app.py

@@ -27,6 +27,8 @@ from html_templates import (
     format_single_dog_result,
     format_multiple_breeds_result,
     format_error_message,
+    format_unknown_breed_message,
+    format_not_dog_message,
     format_warning_html,
     format_multi_dog_container,
     format_breed_details_html,
@@ -238,36 +240,85 @@ def predict_single_dog(image):
             
         return probabilities[0], breeds[:3], relative_probs
         
+# @spaces.GPU
+# def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
+    
+#     results = model_manager.yolo_model(image, conf=conf_threshold, 
+#                                      iou=iou_threshold)[0]
+    
+#     dogs = []
+#     boxes = []
+#     for box in results.boxes:
+#         if box.cls == 16:  # COCO dataset class for dog is 16
+#             xyxy = box.xyxy[0].tolist()
+#             confidence = box.conf.item()
+#             boxes.append((xyxy, confidence))
+
+#     if not boxes:
+#         dogs.append((image, 1.0, [0, 0, image.width, image.height]))
+#     else:
+#         nms_boxes = non_max_suppression(boxes, iou_threshold)
+
+#         for box, confidence in nms_boxes:
+#             x1, y1, x2, y2 = box
+#             w, h = x2 - x1, y2 - y1
+#             x1 = max(0, x1 - w * 0.05)
+#             y1 = max(0, y1 - h * 0.05)
+#             x2 = min(image.width, x2 + w * 0.05)
+#             y2 = min(image.height, y2 + h * 0.05)
+#             cropped_image = image.crop((x1, y1, x2, y2))
+#             dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
+
+#     return dogs
+
 @spaces.GPU
 def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
+    """
+    使用YOLO模型檢測圖片中的狗。
+    只保留被識別為狗（class 16）的物體，並標記它們的狀態。
     
+    Args:
+        image: PIL Image
+        conf_threshold: YOLO檢測的信心度閾值
+        iou_threshold: 非極大值抑制的IoU閾值
+    
+    Returns:
+        list: 包含檢測到的狗的列表，每個元素是(cropped_image, confidence, box, is_dog)的元組
+    """
     results = model_manager.yolo_model(image, conf=conf_threshold, 
                                      iou=iou_threshold)[0]
     
     dogs = []
     boxes = []
+    
+    # 只處理被識別為狗的物體
     for box in results.boxes:
-        if box.cls == 16:  # COCO dataset class for dog is 16
+        class_id = box.cls.item()
+        if class_id == 16:  # COCO dataset中狗的類別是16
             xyxy = box.xyxy[0].tolist()
             confidence = box.conf.item()
-            boxes.append((xyxy, confidence))
+            boxes.append((xyxy, confidence, True))  # 加入is_dog標記
 
     if not boxes:
-        dogs.append((image, 1.0, [0, 0, image.width, image.height]))
-    else:
-        nms_boxes = non_max_suppression(boxes, iou_threshold)
-
-        for box, confidence in nms_boxes:
-            x1, y1, x2, y2 = box
-            w, h = x2 - x1, y2 - y1
-            x1 = max(0, x1 - w * 0.05)
-            y1 = max(0, y1 - h * 0.05)
-            x2 = min(image.width, x2 + w * 0.05)
-            y2 = min(image.height, y2 + h * 0.05)
-            cropped_image = image.crop((x1, y1, x2, y2))
-            dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
-
-    return dogs
+        # 如果沒有檢測到狗，返回整張圖片並標記為非狗
+        return [(image, 1.0, [0, 0, image.width, image.height], False)]
+    
+    nms_boxes = non_max_suppression(boxes, iou_threshold)
+    detected_objects = []
+
+    # 處理每個檢測到的狗
+    for box, confidence, is_dog in nms_boxes:
+        x1, y1, x2, y2 = box
+        w, h = x2 - x1, y2 - y1
+        # 擴大檢測框範圍以包含完整的狗
+        x1 = max(0, x1 - w * 0.05)
+        y1 = max(0, y1 - h * 0.05)
+        x2 = min(image.width, x2 + w * 0.05)
+        y2 = min(image.height, y2 + h * 0.05)
+        cropped_image = image.crop((x1, y1, x2, y2))
+        detected_objects.append((cropped_image, confidence, [x1, y1, x2, y2], is_dog))
+
+    return detected_objects
 
 def non_max_suppression(boxes, iou_threshold):
     keep = []
@@ -324,17 +375,137 @@ def create_breed_comparison(breed1: str, breed2: str) -> dict:
     return comparison_data
 
     
+# def predict(image):
+#     """
+#     Main prediction function that handles both single and multiple dog detection.
+
+#     Args:
+#         image: PIL Image or numpy array
+
+#     Returns:
+#         tuple: (html_output, annotated_image, initial_state)
+#     """
+    
+#     if image is None:
+#         return format_warning_html("Please upload an image to start."), None, None
+
+#     try:
+#         if isinstance(image, np.ndarray):
+#             image = Image.fromarray(image)
+
+#         # Detect dogs in the image
+#         dogs = detect_multiple_dogs(image)
+#         color_scheme = get_color_scheme(len(dogs) == 1)
+
+#         # Prepare for annotation
+#         annotated_image = image.copy()
+#         draw = ImageDraw.Draw(annotated_image)
+
+#         try:
+#             font = ImageFont.truetype("arial.ttf", 24)
+#         except:
+#             font = ImageFont.load_default()
+
+#         dogs_info = ""
+
+#         # Process each detected dog
+#         for i, (cropped_image, detection_confidence, box) in enumerate(dogs):
+#             color = color_scheme if len(dogs) == 1 else color_scheme[i % len(color_scheme)]
+
+#             # Draw box and label on image
+#             draw.rectangle(box, outline=color, width=4)
+#             label = f"Dog {i+1}"
+#             label_bbox = draw.textbbox((0, 0), label, font=font)
+#             label_width = label_bbox[2] - label_bbox[0]
+#             label_height = label_bbox[3] - label_bbox[1]
+
+#             # Draw label background and text
+#             label_x = box[0] + 5
+#             label_y = box[1] + 5
+#             draw.rectangle(
+#                 [label_x - 2, label_y - 2, label_x + label_width + 4, label_y + label_height + 4],
+#                 fill='white',
+#                 outline=color,
+#                 width=2
+#             )
+#             draw.text((label_x, label_y), label, fill=color, font=font)
+
+#             # Predict breed
+#             top1_prob, topk_breeds, relative_probs = predict_single_dog(cropped_image)
+#             combined_confidence = detection_confidence * top1_prob
+
+#             # Format results based on confidence with error handling
+#             try:
+#                 if combined_confidence < 0.2:
+#                     dogs_info += format_error_message(color, i+1)
+#                 elif top1_prob >= 0.45:
+#                     breed = topk_breeds[0]
+#                     description = get_dog_description(breed)
+#                     # Handle missing breed description
+#                     if description is None:
+#                         # 如果沒有描述，創建一個基本描述
+#                         description = {
+#                             "Name": breed,
+#                             "Size": "Unknown",
+#                             "Exercise Needs": "Unknown",
+#                             "Grooming Needs": "Unknown",
+#                             "Care Level": "Unknown",
+#                             "Good with Children": "Unknown",
+#                             "Description": f"Identified as {breed.replace('_', ' ')}"
+#                         }
+#                     dogs_info += format_single_dog_result(breed, description, color)
+#                 else:
+#                     # 修改format_multiple_breeds_result的調用，包含錯誤處理
+#                     dogs_info += format_multiple_breeds_result(
+#                         topk_breeds,
+#                         relative_probs,
+#                         color,
+#                         i+1,
+#                         lambda breed: get_dog_description(breed) or {
+#                             "Name": breed,
+#                             "Size": "Unknown",
+#                             "Exercise Needs": "Unknown",
+#                             "Grooming Needs": "Unknown",
+#                             "Care Level": "Unknown",
+#                             "Good with Children": "Unknown",
+#                             "Description": f"Identified as {breed.replace('_', ' ')}"
+#                         }
+#                     )
+#             except Exception as e:
+#                 print(f"Error formatting results for dog {i+1}: {str(e)}")
+#                 dogs_info += format_error_message(color, i+1)
+
+#         # Wrap final HTML output
+#         html_output = format_multi_dog_container(dogs_info)
+
+#         # Prepare initial state
+#         initial_state = {
+#             "dogs_info": dogs_info,
+#             "image": annotated_image,
+#             "is_multi_dog": len(dogs) > 1,
+#             "html_output": html_output
+#         }
+
+#         return html_output, annotated_image, initial_state
+
+#     except Exception as e:
+#         error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+#         print(error_msg)
+#         return format_warning_html(error_msg), None, None
+
+
+@spaces.GPU
 def predict(image):
     """
-    Main prediction function that handles both single and multiple dog detection.
-
+    主要的預測函數，負責處理狗的檢測和品種辨識。
+    它整合了YOLO的物體檢測和專門的品種分類模型。
+    
     Args:
-        image: PIL Image or numpy array
-
+        image: PIL Image 或 numpy array
+        
     Returns:
         tuple: (html_output, annotated_image, initial_state)
     """
-    
     if image is None:
         return format_warning_html("Please upload an image to start."), None, None
 
@@ -342,11 +513,11 @@ def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # Detect dogs in the image
+        # 檢測圖片中的狗
         dogs = detect_multiple_dogs(image)
         color_scheme = get_color_scheme(len(dogs) == 1)
 
-        # Prepare for annotation
+        # 準備標註
         annotated_image = image.copy()
         draw = ImageDraw.Draw(annotated_image)
 
@@ -357,18 +528,18 @@ def predict(image):
 
         dogs_info = ""
 
-        # Process each detected dog
-        for i, (cropped_image, detection_confidence, box) in enumerate(dogs):
+        # 處理每個檢測到的物體
+        for i, (cropped_image, detection_confidence, box, is_dog) in enumerate(dogs):
             color = color_scheme if len(dogs) == 1 else color_scheme[i % len(color_scheme)]
 
-            # Draw box and label on image
+            # 繪製框和標籤
             draw.rectangle(box, outline=color, width=4)
-            label = f"Dog {i+1}"
+            label = f"Dog {i+1}" if is_dog else f"Object {i+1}"
             label_bbox = draw.textbbox((0, 0), label, font=font)
             label_width = label_bbox[2] - label_bbox[0]
             label_height = label_bbox[3] - label_bbox[1]
 
-            # Draw label background and text
+            # 繪製標籤背景和文字
             label_x = box[0] + 5
             label_y = box[1] + 5
             draw.rectangle(
@@ -379,20 +550,23 @@ def predict(image):
             )
             draw.text((label_x, label_y), label, fill=color, font=font)
 
-            # Predict breed
-            top1_prob, topk_breeds, relative_probs = predict_single_dog(cropped_image)
-            combined_confidence = detection_confidence * top1_prob
-
-            # Format results based on confidence with error handling
             try:
+                # 首先檢查是否為狗
+                if not is_dog:
+                    dogs_info += format_not_dog_message(color, i+1)
+                    continue
+
+                # 如果是狗，進行品種預測
+                top1_prob, topk_breeds, relative_probs = predict_single_dog(cropped_image)
+                combined_confidence = detection_confidence * top1_prob
+
+                # 根據信心度決定輸出格式
                 if combined_confidence < 0.2:
-                    dogs_info += format_error_message(color, i+1)
+                    dogs_info += format_unknown_breed_message(color, i+1)
                 elif top1_prob >= 0.45:
                     breed = topk_breeds[0]
                     description = get_dog_description(breed)
-                    # Handle missing breed description
                     if description is None:
-                        # 如果沒有描述，創建一個基本描述
                         description = {
                             "Name": breed,
                             "Size": "Unknown",
@@ -404,7 +578,6 @@ def predict(image):
                         }
                     dogs_info += format_single_dog_result(breed, description, color)
                 else:
-                    # 修改format_multiple_breeds_result的調用，包含錯誤處理
                     dogs_info += format_multiple_breeds_result(
                         topk_breeds,
                         relative_probs,
@@ -422,12 +595,12 @@ def predict(image):
                     )
             except Exception as e:
                 print(f"Error formatting results for dog {i+1}: {str(e)}")
-                dogs_info += format_error_message(color, i+1)
+                dogs_info += format_unknown_breed_message(color, i+1)
 
-        # Wrap final HTML output
+        # 包裝最終的HTML輸出
         html_output = format_multi_dog_container(dogs_info)
 
-        # Prepare initial state
+        # 準備初始狀態
         initial_state = {
             "dogs_info": dogs_info,
             "image": annotated_image,