Update app.py

app.py

@@ -312,7 +312,7 @@ def _predict_single_dog(image):
 #     return dogs
 # 此為如果後面調不好 使用的版本
 
-async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.5):
+async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.3):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     for box in results.boxes:
@@ -320,42 +320,31 @@ async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.5):
             xyxy = box.xyxy[0].tolist()
             confidence = box.conf.item()
             area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
-            if area > 1000:  # 過濾掉太小的檢測框
+            image_area = image.width * image.height
+            if area > 0.01 * image_area:  # 過濾掉太小的檢測框，但使用相對面積
                 cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
                 dogs.append((cropped_image, confidence, xyxy))
     
-    # 合併重疊的檢測框
-    dogs = merge_overlapping_boxes(dogs, iou_threshold=0.6)
+    # 如果檢測到的狗太少，嘗試降低閾值再次檢測
+    if len(dogs) < 2:
+        results = model_yolo(image, conf=conf_threshold/2, iou=iou_threshold)[0]
+        for box in results.boxes:
+            if box.cls == 16:
+                xyxy = box.xyxy[0].tolist()
+                confidence = box.conf.item()
+                area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
+                image_area = image.width * image.height
+                if area > 0.01 * image_area and not is_box_duplicate(xyxy, [d[2] for d in dogs]):
+                    cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
+                    dogs.append((cropped_image, confidence, xyxy))
     
     return dogs
 
-def merge_overlapping_boxes(dogs, iou_threshold=0.6):
-    merged_dogs = []
-    while dogs:
-        base = dogs.pop(0)
-        i = 0
-        while i < len(dogs):
-            if calculate_iou(base[2], dogs[i][2]) > iou_threshold:
-                # 合併重疊的框
-                base = merge_boxes(base, dogs.pop(i))
-            else:
-                i += 1
-        merged_dogs.append(base)
-    return merged_dogs
-
-def merge_boxes(box1, box2):
-    xyxy1, conf1, _ = box1
-    xyxy2, conf2, _ = box2
-    merged_xyxy = [
-        min(xyxy1[0], xyxy2[0]),
-        min(xyxy1[1], xyxy2[1]),
-        max(xyxy1[2], xyxy2[2]),
-        max(xyxy1[3], xyxy2[3])
-    ]
-    merged_conf = max(conf1, conf2)
-    merged_image = Image.new('RGB', (int(merged_xyxy[2] - merged_xyxy[0]), int(merged_xyxy[3] - merged_xyxy[1])))
-    merged_image.paste(box1[0], (0, 0))
-    return (merged_image, merged_conf, merged_xyxy)
+def is_box_duplicate(new_box, existing_boxes, iou_threshold=0.5):
+    for box in existing_boxes:
+        if calculate_iou(new_box, box) > iou_threshold:
+            return True
+    return False
 
 def calculate_iou(box1, box2):
     # 計算兩個邊界框的交集面積
@@ -494,15 +483,15 @@ async def predict(image):
             image = Image.fromarray(image)
 
         dogs = await detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.5)
-        
-        # 如果檢測到的狗的數量不合理，嘗試調整參數重新檢測
-        if len(dogs) > 5 or (len(dogs) == 0 and has_dog_features(image)):
-            dogs = await detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.4)
-        
+          
         if len(dogs) == 0:
             return await process_single_dog(image)
         elif len(dogs) == 1:
-            return await process_single_dog(dogs[0][0])
+            # 如果只檢測到一隻狗，但圖像可能包含多隻狗，再次嘗試檢測
+            if has_multiple_dogs(image):
+                dogs = await detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.2)
+            if len(dogs) == 1:
+                return await process_single_dog(dogs[0][0])
         else:
              # 多狗情境
             color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
@@ -558,12 +547,13 @@ async def predict(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
 
-def has_dog_features(image):
-    # 使用簡單的啟發式方法來檢查圖像是否可能包含狗
+def has_multiple_dogs(image):
+    # 使用簡單的啟發式方法來檢查圖像是否可能包含多隻狗
     # 這裡可以使用更複雜的方法，如特徵提取或輕量級模型
     gray = image.convert('L')
     edges = gray.filter(ImageFilter.FIND_EDGES)
-    return np.mean(np.array(edges)) > 10  # 假設邊緣檢測後的平均值大於 10 表示可能有狗
+    edge_pixels = np.array(edges)
+    return np.sum(edge_pixels > 128) > image.width * image.height * 0.1  # 假設邊緣像素比例大於 10% 表示可能有多隻狗
 
 async def process_single_dog(image):
     top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)