Update app.py

app.py

@@ -243,14 +243,12 @@ def _predict_single_dog(image):
 #         print(error_msg)  # 添加日誌輸出
 #         return error_msg, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
 
-async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.5, merge_threshold=0.2):
+async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.4, merge_threshold=0.3):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
-    boxes = []
-    confidences = []
     
     image_area = image.width * image.height
-    min_area_ratio = 0.005  # 最小檢測面積佔整個圖像的比例
+    min_area_ratio = 0.01  # 最小檢測面積佔整個圖像的比例
     
     for box in results.boxes:
         if box.cls == 16:  # COCO 數據集中狗的類別是 16
@@ -258,72 +256,34 @@ async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.5, mer
             area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
             if area / image_area >= min_area_ratio:
                 confidence = box.conf.item()
-                boxes.append(torch.tensor(xyxy))
-                confidences.append(confidence)
+                cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
+                dogs.append((cropped_image, confidence, xyxy))
     
-    if boxes:
-        boxes = torch.stack(boxes)
-        confidences = torch.tensor(confidences)
+    # 合併重疊的框
+    merged_dogs = []
+    while dogs:
+        base_dog = dogs.pop(0)
+        base_box = torch.tensor(base_dog[2])
+        to_merge = [base_dog]
         
-        # 應用軟 NMS
-        keep = soft_nms(boxes, confidences, iou_threshold=iou_threshold, sigma=0.5)
-        
-        for i in keep:
-            xyxy = boxes[i].tolist()
-            confidence = confidences[i].item()
-            cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
-            dogs.append((cropped_image, confidence, xyxy))
-        
-        # 合併鄰近的邊界框
-        merged_dogs = []
-        while dogs:
-            base_dog = dogs.pop(0)
-            base_box = torch.tensor(base_dog[2])
-            to_merge = [base_dog]
-            
-            i = 0
-            while i < len(dogs):
-                compare_box = torch.tensor(dogs[i][2])
-                iou = box_iou(base_box.unsqueeze(0), compare_box.unsqueeze(0)).item()
-                if iou > merge_threshold:
-                    to_merge.append(dogs.pop(i))
-                else:
-                    i += 1
-            
-            if len(to_merge) == 1:
-                merged_dogs.append(base_dog)
+        i = 0
+        while i < len(dogs):
+            compare_box = torch.tensor(dogs[i][2])
+            iou = box_iou(base_box.unsqueeze(0), compare_box.unsqueeze(0)).item()
+            if iou > merge_threshold:
+                to_merge.append(dogs.pop(i))
             else:
-                merged_box = torch.cat([torch.tensor(dog[2]).unsqueeze(0) for dog in to_merge]).mean(0)
-                merged_confidence = max(dog[1] for dog in to_merge)
-                merged_image = image.crop(merged_box.tolist())
-                merged_dogs.append((merged_image, merged_confidence, merged_box.tolist()))
-        
-        # 後處理：限制檢測到的狗狗數量
-        if len(merged_dogs) > 5:
-            merged_dogs = sorted(merged_dogs, key=lambda x: x[1], reverse=True)[:5]
+                i += 1
         
-        return merged_dogs
-    
-    # 如果沒有檢測到狗狗，使用備用分類器
-    return await fallback_classifier(image)
-
-async def fallback_classifier(image):
-    # 使用預訓練的 ResNet 或其他適合的分類器
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
-    img_tensor = transform(image).unsqueeze(0)
-    
-    with torch.no_grad():
-        output = fallback_model(img_tensor)
-        confidence, predicted = torch.max(output, 1)
+        if len(to_merge) == 1:
+            merged_dogs.append(base_dog)
+        else:
+            merged_box = torch.tensor([dog[2] for dog in to_merge]).mean(0)
+            merged_confidence = max(dog[1] for dog in to_merge)
+            merged_image = image.crop(merged_box.tolist())
+            merged_dogs.append((merged_image, merged_confidence, merged_box.tolist()))
     
-    if confidence.item() > 0.5:  # 設置一個合適的閾值
-        return [(image, confidence.item(), [0, 0, image.width, image.height])]
-    else:
-        return []
+    return merged_dogs if merged_dogs else [(image, 1.0, [0, 0, image.width, image.height])]
 
 async def predict(image):
     if image is None:
@@ -333,15 +293,12 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        dogs = await detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4, merge_threshold=0.3)
-        
-        if len(dogs) == 0:
-            return "No dogs detected in the image.", image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
+        dogs = await detect_multiple_dogs(image)
         
         if len(dogs) == 1:
-            return await process_single_dog(dogs[0][0])  # Pass the cropped image of the single detected dog
+            return await process_single_dog(dogs[0][0])
 
-        # Multi-dog scenario
+        # 多狗情境
         color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
         explanations = []
         buttons = []
@@ -349,7 +306,7 @@ async def predict(image):
         draw = ImageDraw.Draw(annotated_image)
         font = ImageFont.load_default()
 
-        for i, (cropped_image, confidence, box) in enumerate(dogs):
+        for i, (cropped_image, _, 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)
@@ -360,13 +317,11 @@ async def predict(image):
                 description = get_dog_description(breed)
                 formatted_description = format_description(description, breed)
                 explanations.append(f"Dog {i+1}: {formatted_description}")
-            elif top1_prob >= 0.2:
+            else:
                 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([gr.update(visible=True, value=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.")
 
         final_explanation = "\n\n".join(explanations)
         if buttons:
@@ -392,7 +347,7 @@ async def predict(image):
 
     except Exception as e:
         error_msg = f"An error occurred: {str(e)}"
-        print(error_msg)  # Add log output
+        print(error_msg)  # 添加日誌輸出
         return error_msg, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None