Update app.py

app.py

@@ -9,6 +9,12 @@ from torchvision import transforms
 from PIL import Image
 from data_manager import get_dog_description
 from urllib.parse import quote
+os.system('pip install ultralytics')
+from ultralytics import YOLO
+
+# 下載YOLOv5預訓練模型
+model_yolo = YOLO('yolov5s.pt')  # 使用 YOLOv5 預訓練模型
+
 
 dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staffordshire_Terrier", 
               "Appenzeller", "Australian_Terrier", "Bedlington_Terrier", "Bernese_Mountain_Dog", 
@@ -121,44 +127,101 @@ def preprocess_image(image):
 def get_akc_breeds_link():
     return "https://www.akc.org/dog-breeds/"
 
+# def predict(image):
+#     if image is None:
+#         return "Please upload an image to get started.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+    
+#     try:
+#         image_tensor = preprocess_image(image)
+#         with torch.no_grad():
+#             output = model(image_tensor)
+#             logits = output[0] if isinstance(output, tuple) else output
+
+#             probabilities = F.softmax(logits, dim=1)
+#             topk_probs, topk_indices = torch.topk(probabilities, k=3)
+
+#             top1_prob = topk_probs[0][0].item()
+#             topk_breeds = [dog_breeds[idx.item()] for idx in topk_indices[0]]
+#             topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
+
+#             if top1_prob >= 0.5:
+#                 breed = topk_breeds[0]
+#                 description = get_dog_description(breed)
+#                 return format_description(description, breed), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+#             elif top1_prob < 0.2:
+#                 return ("The image is too unclear or the dog breed is not in the dataset. Please upload a clearer image of the dog.",
+#                         gr.update(visible=False), gr.update(visible=False), gr.update(visible=False))
+#             else:
+#                 explanation = (
+#                     f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n"
+#                     f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]} confidence)\n"
+#                     f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]} confidence)\n"
+#                     f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n\n"
+#                     "Click on a button to view more information about the breed."
+#                 )
+#                 return explanation, gr.update(visible=True, value=f"More about {topk_breeds[0]}"), gr.update(visible=True, value=f"More about {topk_breeds[1]}"), gr.update(visible=True, value=f"More about {topk_breeds[2]}")
+
+#     except Exception as e:
+#         return f"An error occurred: {e}", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
 def predict(image):
     if image is None:
         return "Please upload an image to get started.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-    
+
     try:
-        image_tensor = preprocess_image(image)
-        with torch.no_grad():
-            output = model(image_tensor)
-            logits = output[0] if isinstance(output, tuple) else output
-
-            probabilities = F.softmax(logits, dim=1)
-            topk_probs, topk_indices = torch.topk(probabilities, k=3)
-
-            top1_prob = topk_probs[0][0].item()
-            topk_breeds = [dog_breeds[idx.item()] for idx in topk_indices[0]]
-            topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
-
-            if top1_prob >= 0.5:
-                breed = topk_breeds[0]
-                description = get_dog_description(breed)
-                return format_description(description, breed), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-
-            elif top1_prob < 0.2:
-                return ("The image is too unclear or the dog breed is not in the dataset. Please upload a clearer image of the dog.",
-                        gr.update(visible=False), gr.update(visible=False), gr.update(visible=False))
-            else:
-                explanation = (
-                    f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n"
-                    f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]} confidence)\n"
-                    f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]} confidence)\n"
-                    f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n\n"
-                    "Click on a button to view more information about the breed."
-                )
-                return explanation, gr.update(visible=True, value=f"More about {topk_breeds[0]}"), gr.update(visible=True, value=f"More about {topk_breeds[1]}"), gr.update(visible=True, value=f"More about {topk_breeds[2]}")
+        # 使用 YOLO 偵測狗
+        results = model_yolo(image)
+        dogs = results.xyxy[0]  # 提取偵測到的狗的邊界框
+
+        if len(dogs) == 0:
+            return "No dog detected in the image.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+        explanations = []
+        visible_buttons = []
+
+        for i, box in enumerate(dogs):
+            x1, y1, x2, y2 = map(int, box[:4])
+            cropped_image = image.crop((x1, y1, x2, y2))  # 裁剪狗區域
+            image_tensor = preprocess_image(cropped_image)
+
+            with torch.no_grad():
+                output = model(image_tensor)
+                logits = output[0] if isinstance(output, tuple) else output
+
+                probabilities = F.softmax(logits, dim=1)
+                topk_probs, topk_indices = torch.topk(probabilities, k=3)
+
+                top1_prob = topk_probs[0][0].item()
+                topk_breeds = [dog_breeds[idx.item()] for idx in topk_indices[0]]
+                topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
+
+                # 根據信心分數進行判斷
+                if top1_prob >= 0.5:
+                    breed = topk_breeds[0]
+                    description = get_dog_description(breed)
+                    explanations.append(f"Detected a dog: **{breed}** with {topk_probs_percent[0]} confidence.")
+                elif 0.2 <= top1_prob < 0.5:
+                    explanation = (
+                        f"Detected a dog with moderate confidence. Here are the top 3 possible breeds:\n"
+                        f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]} confidence)\n"
+                        f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]} confidence)\n"
+                        f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n"
+                    )
+                    explanations.append(explanation)
+                    visible_buttons.extend([i+1 for _ in range(3)])
+                else:
+                    explanations.append("The image is too unclear or the breed is not in the dataset. Please upload a clearer image.")
+
+        # 處理不同情境的結果
+        if len(explanations) > 0:
+            final_explanation = "\n\n".join(explanations)
+            return final_explanation, gr.update(visible=len(visible_buttons) >= 1), gr.update(visible=len(visible_buttons) >= 2), gr.update(visible=len(visible_buttons) >= 3)
 
     except Exception as e:
         return f"An error occurred: {e}", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
+
 def format_description(description, breed):
     if isinstance(description, dict):
         formatted_description = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])