Update app.py

app.py

@@ -311,49 +311,63 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # YOLO 偵測多隻狗
+        # 使用 YOLO 偵測有無多隻狗
         dogs = await detect_multiple_dogs(image)
+
+        if len(dogs) == 1:
+            # 如果只有一隻狗，直接預測品種，不使用 YOLO
+            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
+            breed = topk_breeds[0]
+            description = get_dog_description(breed)
+            formatted_description = format_description(description, breed)
+            
+            if top1_prob < 0.2:
+                return "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+            
+            return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
         
-        if len(dogs) == 0:
-            return "No dogs detected. Please upload a clear image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-
-        explanations = []
-        buttons = []
-        annotated_image = image.copy()
-        draw = ImageDraw.Draw(annotated_image)
-        font = ImageFont.load_default()
-
-        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)
-
-            if top1_prob >= 0.5:
-                breed = topk_breeds[0]
-                description = get_dog_description(breed)
-                explanations.append(f"Dog {i}: **{breed}**\n\n{format_description(description, breed, is_multi_dog=True, dog_number=i)}")
-            else:
-                explanation = f"""
+        elif len(dogs) > 1:
+            # 如果偵測到多隻狗，使用 YOLO 進行分割並預測
+            explanations = []
+            buttons = []
+            annotated_image = image.copy()
+            draw = ImageDraw.Draw(annotated_image)
+            font = ImageFont.load_default()
+
+            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)
+
+                if top1_prob >= 0.5:
+                    breed = topk_breeds[0]
+                    description = get_dog_description(breed)
+                    explanations.append(f"Dog {i}: **{breed}**\n\n{format_description(description, breed, is_multi_dog=True, dog_number=i)}")
+                else:
+                    explanation = f"""
 Dog {i}: Detected with moderate confidence. Here are the top 3 possible breeds:
 1. **{topk_breeds[0]}** ({topk_probs_percent[0]})
 2. **{topk_breeds[1]}** ({topk_probs_percent[1]})
 3. **{topk_breeds[2]}** ({topk_probs_percent[2]})
 """
-                explanations.append(explanation)
-                for breed in topk_breeds:
-                    buttons.append(f"More about Dog {i}: {breed}")
+                    explanations.append(explanation)
+                    for breed in topk_breeds:
+                        buttons.append(f"More about Dog {i}: {breed}")
+
+            final_explanation = "\n\n---\n\n".join(explanations)
+            
+            if buttons:
+                return final_explanation, annotated_image, gr.update(visible=True, choices=buttons), gr.update(visible=False), gr.update(visible=False)
+            else:
+                return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
-        final_explanation = "\n\n---\n\n".join(explanations)
-        
-        if buttons:
-            return final_explanation, annotated_image, gr.update(visible=True, choices=buttons), gr.update(visible=False), gr.update(visible=False)
         else:
-            return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-
+            # 未偵測到狗
+            return "No dogs detected. Please upload a clear image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible(False))
+            
     except Exception as e:
-        return f"An error occurred: {e}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+        return f"An error occurred: {e}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible(False))
 
 # 顯示選擇的品種詳細信息
 async def show_details(choice):