Update app.py

app.py

@@ -283,14 +283,13 @@ async def predict_single_dog(image):
         topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
     return top1_prob, topk_breeds, topk_probs_percent
 
-# 偵測多隻狗的函數
+
 async def detect_multiple_dogs(image):
     try:
-        # 調整圖像大小以加快處理速度
         img = image.copy()
         img.thumbnail((640, 640))
         
-        results = model_yolo(img, conf=0.3)  # 降低置信度閾值以檢測更多狗
+        results = model_yolo(img, conf=0.2)  # 降低閾值以檢測更多狗
         dogs = []
         for result in results:
             for box in result.boxes:
@@ -299,6 +298,15 @@ async def detect_multiple_dogs(image):
                     confidence = box.conf.item()
                     cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
                     dogs.append((cropped_image, confidence, xyxy))
+        
+        # 如果只檢測到一隻狗，嘗試檢測其他可能的狗
+        if len(dogs) == 1:
+            # 使用整張圖像進行品種預測
+            full_image_prob, full_image_breeds, _ = await predict_single_dog(image)
+            if full_image_prob >= 0.3 and full_image_breeds[0] != dogs[0][0]:
+                # 如果整張圖像的預測結果不同且置信度較高，添加為第二隻狗
+                dogs.append((image, full_image_prob, [0, 0, image.width, image.height]))
+        
         return dogs
     except Exception as e:
         print(f"Error in detect_multiple_dogs: {e}")
@@ -312,11 +320,15 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # 使用 YOLO 檢測狗
         dogs = await detect_multiple_dogs(image)
 
         if len(dogs) == 0:
-            return "No dogs detected in the image. Please upload a clear image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+            # 如果沒有檢測到狗，嘗試對整張圖像進行預測
+            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
+            if top1_prob >= 0.3:
+                return await process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent, image, [0, 0, image.width, image.height])
+            else:
+                return "No dogs detected in the image. Please upload a clear image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
         elif len(dogs) == 1:
             cropped_image, _, box = dogs[0]
             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
@@ -338,8 +350,9 @@ async def process_multiple_dogs_result(dogs, image):
     for i, (cropped_image, _, box) in enumerate(dogs, 1):
         top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
         
-        draw.rectangle(box, outline="red", width=3)
-        draw.text((box[0], box[1]), f"Dog {i}", fill="yellow", font=font)
+        optimized_box = optimize_box(box, image.size)
+        draw.rectangle(optimized_box, outline="red", width=3)
+        draw.text((optimized_box[0], optimized_box[1]), f"Dog {i}", fill="yellow", font=font)
         
         if top1_prob >= 0.5:
             breed = topk_breeds[0]
@@ -366,8 +379,9 @@ Dog {i}: Detected with moderate confidence. Here are the top 3 possible breeds:
 async def process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent, image, box):
     annotated_image = image.copy()
     draw = ImageDraw.Draw(annotated_image)
-    draw.rectangle(box, outline="red", width=3)
-    draw.text((box[0], box[1]), "Dog", fill="yellow", font=ImageFont.load_default())
+    optimized_box = optimize_box(box, image.size)
+    draw.rectangle(optimized_box, outline="red", width=3)
+    draw.text((optimized_box[0], optimized_box[1]), "Dog", fill="yellow", font=ImageFont.load_default())
 
     if top1_prob >= 0.5:
         breed = topk_breeds[0]
@@ -387,6 +401,16 @@ async def process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent,
     else:
         return "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.", annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
+def optimize_box(box, image_size):
+    x1, y1, x2, y2 = box
+    w, h = image_size
+    # 擴大邊界框以確保完整包含狗
+    x1 = max(0, x1 - 10)
+    y1 = max(0, y1 - 10)
+    x2 = min(w, x2 + 10)
+    y2 = min(h, y2 + 10)
+    return [x1, y1, x2, y2]
+
 async def show_details(choice):
     if not choice:
         return "Please select a breed to view details."