Update app.py

app.py

@@ -10,6 +10,7 @@ from PIL import Image, ImageDraw, ImageFont
 from data_manager import get_dog_description
 from urllib.parse import quote
 from ultralytics import YOLO
+import asyncio
 
 
 # 下載YOLOv8預訓練模型
@@ -270,7 +271,19 @@ Please refer to the AKC's terms of use and privacy policy.*
 """
     return formatted_description
 
-def predict_single_dog(image):
+# 全局變量，用於存儲加載的模型
+model = None
+model_yolo = None
+
+def load_models():
+    global model, model_yolo
+    model = # 加載您的品種分類模型
+    model_yolo = YOLO('yolov8n.pt')
+
+# 在應用啟動時加載模型
+load_models()
+
+async def predict_single_dog(image):
     image_tensor = preprocess_image(image)
     with torch.no_grad():
         output = model(image_tensor)
@@ -282,7 +295,7 @@ 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):
+async def detect_multiple_dogs(image):
     try:
         results = model_yolo(image)
         dogs = []
@@ -298,7 +311,7 @@ def detect_multiple_dogs(image):
         print(f"Error in detect_multiple_dogs: {e}")
         return []
 
-def predict(image):
+async def predict(image):
     if image is None:
         return "Please upload an image to start.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
@@ -306,13 +319,11 @@ def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # Always start with single dog prediction
-        top1_prob, topk_breeds, topk_probs_percent = predict_single_dog(image)
-        
-        # Check if we need to use YOLO for multiple dogs
-        dogs = detect_multiple_dogs(image)
+        # 快速檢查是否有多個狗
+        dogs = await detect_multiple_dogs(image)
         
-        if len(dogs) <= 1:  # Single dog or no dog detected
+        if len(dogs) <= 1:  # 單狗或無狗情況
+            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)
@@ -333,7 +344,7 @@ Click on a button below to view more information about each breed.
             else:
                 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)
         
-        # Multiple dogs detected, process each dog
+        # 多狗情況
         explanations = []
         visible_buttons = []
         annotated_image = image.copy()
@@ -341,7 +352,7 @@ Click on a button below to view more information about each breed.
         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)
+            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="red", font=font)