Update app.py

app.py

@@ -167,33 +167,6 @@ 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.4):
-#     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.25, iou_threshold=0.4):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
@@ -212,19 +185,16 @@ async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4):
         for box, confidence in nms_boxes:
             x1, y1, x2, y2 = box
             w, h = x2 - x1, y2 - y1
-            x1 = max(0, x1 - w * 0.1)
-            y1 = max(0, y1 - h * 0.1)
-            x2 = min(image.width, x2 + w * 0.1)
-            y2 = min(image.height, y2 + h * 0.1)
+            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]))
     
-    # 如果只检测到一只狗，但置信度较低，添加整张图片作为备选
-    if len(dogs) == 1 and dogs[0][1] < 0.5:
-        dogs.append((image, 1.0, [0, 0, image.width, image.height]))
-    
     return dogs
 
+
 def non_max_suppression(boxes, iou_threshold):
     keep = []
     boxes = sorted(boxes, key=lambda x: x[1], reverse=True)
@@ -296,72 +266,6 @@ 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
@@ -383,9 +287,7 @@ async def predict(image):
             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)
-            
-            if len(dogs) > 1:
-                draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
+            draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
 
             combined_confidence = detection_confidence * top1_prob
 
@@ -393,23 +295,14 @@ async def predict(image):
                 breed = topk_breeds[0]
                 description = get_dog_description(breed)
                 formatted_description = format_description(description, breed)
-                if len(dogs) == 1:
-                    explanations.append(f"Breed: {breed}\n{formatted_description}")
-                else:
-                    explanations.append(f"Dog {i+1}: Breed: {breed}\n{formatted_description}")
+                explanations.append(f"Dog {i+1}: {formatted_description}")
             elif combined_confidence >= 0.2:
-                if len(dogs) == 1:
-                    dog_explanation = f"Top 3 possible breeds:\n"
-                else:
-                    dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
+                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 ' + str(i+1) + ': ' if len(dogs) > 1 else ''}More about {breed}" for breed in topk_breeds[:3]])
+                buttons.extend([f"Dog {i+1}: More about {breed}" for breed in topk_breeds[:3]])
             else:
-                if len(dogs) == 1:
-                    explanations.append("The image is unclear or the breed is not in the dataset. Please upload a clearer image.")
-                else:
-                    explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset. Please upload a clearer image.")
+                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:
@@ -440,6 +333,7 @@ async def predict(image):
         return error_msg, None, gr.update(visible=False, choices=[]), None
 
 
+
 def show_details(choice, previous_output, initial_state):
     if not choice:
         return previous_output, gr.update(visible=True), initial_state