Update app.py

app.py

@@ -167,34 +167,56 @@ async def predict_single_dog(image):
     return top1_prob, topk_breeds, topk_probs_percent
     
 
-async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.6):
+# async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.6):
+#     results = model_yolo(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
+
+async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.5):
     results = model_yolo(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
+        if box.cls == 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
 
 
+
 def non_max_suppression(boxes, iou_threshold):
     keep = []
     boxes = sorted(boxes, key=lambda x: x[1], reverse=True)
@@ -266,9 +288,76 @@ async def process_single_dog(image):
         return explanation, image, buttons[0], buttons[1], buttons[2], gr.update(visible=True), initial_state
 
 
+# async def predict(image):
+#     if image is None:
+#         return "Please upload an image to start.", None, gr.update(visible=False, choices=[]), None
+
+#     try:
+#         if isinstance(image, np.ndarray):
+#             image = Image.fromarray(image)
+
+#         dogs = await detect_multiple_dogs(image)
+        
+#         color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
+#         explanations = []
+#         buttons = []
+#         annotated_image = image.copy()
+#         draw = ImageDraw.Draw(annotated_image)
+#         font = ImageFont.load_default()
+
+#         for i, (cropped_image, detection_confidence, box) in enumerate(dogs):
+#             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
+#             color = color_list[i % len(color_list)]
+#             draw.rectangle(box, outline=color, width=3)
+#             draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
+
+#             combined_confidence = detection_confidence * top1_prob
+
+#             if top1_prob >= 0.5:
+#                 breed = topk_breeds[0]
+#                 description = get_dog_description(breed)
+#                 formatted_description = format_description(description, breed)
+#                 explanations.append(f"Dog {i+1}: {formatted_description}")
+#             elif combined_confidence >= 0.2:
+#                 dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
+#                 dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
+#                 explanations.append(dog_explanation)
+#                 buttons.extend([f"Dog {i+1}: More about {breed}" for breed in topk_breeds[:3]])
+#             else:
+#                 explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset. Please upload a clearer image.")
+
+#         final_explanation = "\n\n".join(explanations)
+#         if buttons:
+#             final_explanation += "\n\nClick on a button to view more information about the breed."
+#             initial_state = {
+#                 "explanation": final_explanation,
+#                 "buttons": buttons,
+#                 "show_back": True,
+#                 "image": annotated_image,
+#                 "is_multi_dog": len(dogs) > 1,
+#                 "dogs_info": explanations
+#             }
+#             return final_explanation, annotated_image, gr.update(visible=True, choices=buttons), initial_state
+#         else:
+#             initial_state = {
+#                 "explanation": final_explanation,
+#                 "buttons": [],
+#                 "show_back": False,
+#                 "image": annotated_image,
+#                 "is_multi_dog": len(dogs) > 1,
+#                 "dogs_info": explanations
+#             }
+#             return final_explanation, annotated_image, gr.update(visible=False, choices=[]), initial_state
+
+#     except Exception as e:
+#         error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+#         print(error_msg)
+#         return error_msg, None, gr.update(visible=False, choices=[]), None
+
+
 async def predict(image):
     if image is None:
-        return "Please upload an image to start.", None, gr.update(visible=False, choices=[]), None
+        return "Please upload an image to start.", None, gr.update(visible=False), None
 
     try:
         if isinstance(image, np.ndarray):
@@ -276,62 +365,30 @@ async def predict(image):
 
         dogs = await detect_multiple_dogs(image)
         
-        color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
         explanations = []
-        buttons = []
         annotated_image = image.copy()
         draw = ImageDraw.Draw(annotated_image)
-        font = ImageFont.load_default()
-
+        
+        # 針對每隻狗進行辨識
         for i, (cropped_image, detection_confidence, box) in enumerate(dogs):
             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
-            color = color_list[i % len(color_list)]
-            draw.rectangle(box, outline=color, width=3)
-            draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
-
+            draw.rectangle(box, outline="#FF0000", width=3)  # 標記框選
             combined_confidence = detection_confidence * top1_prob
 
-            if top1_prob >= 0.5:
+            if combined_confidence >= 0.5:
                 breed = topk_breeds[0]
                 description = get_dog_description(breed)
-                formatted_description = format_description(description, breed)
-                explanations.append(f"Dog {i+1}: {formatted_description}")
-            elif combined_confidence >= 0.2:
-                dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
-                dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
-                explanations.append(dog_explanation)
-                buttons.extend([f"Dog {i+1}: More about {breed}" for breed in topk_breeds[:3]])
+                explanations.append(f"Dog {i+1}: {description}")
             else:
-                explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset. Please upload a clearer image.")
-
-        final_explanation = "\n\n".join(explanations)
-        if buttons:
-            final_explanation += "\n\nClick on a button to view more information about the breed."
-            initial_state = {
-                "explanation": final_explanation,
-                "buttons": buttons,
-                "show_back": True,
-                "image": annotated_image,
-                "is_multi_dog": len(dogs) > 1,
-                "dogs_info": explanations
-            }
-            return final_explanation, annotated_image, gr.update(visible=True, choices=buttons), initial_state
-        else:
-            initial_state = {
-                "explanation": final_explanation,
-                "buttons": [],
-                "show_back": False,
-                "image": annotated_image,
-                "is_multi_dog": len(dogs) > 1,
-                "dogs_info": explanations
-            }
-            return final_explanation, annotated_image, gr.update(visible=False, choices=[]), initial_state
+                explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
+        
+        # 返回最終解釋與圖片
+        return "\n".join(explanations), annotated_image
 
     except Exception as e:
-        error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+        error_msg = f"Error: {str(e)}"
         print(error_msg)
-        return error_msg, None, gr.update(visible=False, choices=[]), None
-
+        return error_msg, None
 
 
 def show_details(choice, previous_output, initial_state):