Update app.py

app.py

@@ -6,12 +6,10 @@ import gradio as gr
 from torchvision.models import efficientnet_v2_m, EfficientNet_V2_M_Weights
 import torch.nn.functional as F
 from torchvision import transforms
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 from data_manager import get_dog_description
 from urllib.parse import quote
-# os.system('pip install ultralytics')
 from ultralytics import YOLO
-from PIL import ImageDraw
 
 
 # 下載YOLOv8預訓練模型
@@ -273,7 +271,6 @@ Please refer to the AKC's terms of use and privacy policy.*
     return formatted_description
 
 def predict_single_dog(image):
-    # 直接使用模型進行預測，無需通過 YOLO
     image_tensor = preprocess_image(image)
     with torch.no_grad():
         output = model(image_tensor)
@@ -285,20 +282,21 @@ 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
 
-
 def detect_multiple_dogs(image):
-    # 使用 YOLO 檢測多隻狗
-    results = model_yolo(image)
-    dogs = []
-    for result in results:
-        for box in result.boxes:
-            if box.cls == 16:  # COCO 資料集中狗的類別是 16
-                xyxy = box.xyxy[0].tolist()
-                confidence = box.conf.item()
-                cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
-                dogs.append((cropped_image, confidence, xyxy))
-    return dogs
-
+    try:
+        results = model_yolo(image)
+        dogs = []
+        for result in results:
+            for box in result.boxes:
+                if box.cls == 16:  # COCO dataset class for dog is 16
+                    xyxy = box.xyxy[0].tolist()
+                    confidence = box.conf.item()
+                    cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
+                    dogs.append((cropped_image, confidence, xyxy))
+        return dogs
+    except Exception as e:
+        print(f"Error in detect_multiple_dogs: {e}")
+        return []
 
 def predict(image):
     if image is None:
@@ -340,12 +338,13 @@ Click on a button below to view more information about each breed.
         visible_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 = predict_single_dog(cropped_image)
 
             draw.rectangle(box, outline="red", width=3)
-            draw.text((box[0], box[1]), f"Dog {i}", fill="red", font=ImageFont.truetype("arial.ttf", 20))
+            draw.text((box[0], box[1]), f"Dog {i}", fill="red", font=font)
 
             if top1_prob >= 0.5:
                 breed = topk_breeds[0]