Update app.py

app.py

@@ -35,6 +35,7 @@ from urllib.parse import quote
 from ultralytics import YOLO
 import traceback
 import spaces
+import asyncio
 
 # model_yolo = YOLO('yolov8l.pt')
 
@@ -601,55 +602,74 @@ def preprocess_image(image):
     ])
 
     return transform(image).unsqueeze(0)
+
+global_models = {
+    'classification_model': None,
+    'detection_model': None
+}
+
+@spaces.GPU
+async def initialize_models():
+    """
+    初始化並預熱所有模型
+    這個函數集中管理所有模型的初始化，確保它們只被載入一次
+    """
+    try:
+        # 初始化分類模型
+        if global_models['classification_model'] is None:
+            model = BaseModel(num_classes=len(dog_breeds), device='cuda').to('cuda')
+            checkpoint = torch.load('124_best_model_dog.pth', map_location='cuda')
+            model.load_state_dict(checkpoint['base_model'], strict=False)
+            model.eval()
+            global_models['classification_model'] = model
+        
+        # 初始化檢測模型
+        if global_models['detection_model'] is None:
+            global_models['detection_model'] = YOLO('yolov8l.pt')
+            
+        return True
+    except Exception as e:
+        print(f"模型初始化失敗: {str(e)}")
+        return False
+
     
 @spaces.GPU
 async def predict_single_dog(image):
     """
-    Predicts the dog breed using only the classifier.
-    Args:
-        image: PIL Image or numpy array
-    Returns:
-        tuple: (top1_prob, topk_breeds, relative_probs)
+    預測狗狗品種
+    使用全域模型進行預測，避免重複載入
     """
-    if not hasattr(predict_single_dog, 'model'):
-        num_classes = len(dog_breeds)
-        predict_single_dog.model = BaseModel(num_classes=len(dog_breeds), device='cuda').to('cuda')
-        model_path = '124_best_model_dog.pth'
-        checkpoint = torch.load(model_path, map_location='cuda')
-        predict_single_dog.model.load_state_dict(checkpoint['base_model'], strict=False)
-        predict_single_dog.model.eval()
+    model = global_models['classification_model']
+    if model is None:
+        await initialize_models()
+        model = global_models['classification_model']
     
     image_tensor = preprocess_image(image).to('cuda')
     
     with torch.no_grad():
-        # Get model outputs (只使用logits，不需要features)
-        logits = predict_single_dog.model(image_tensor)[0]
+        logits = model(image_tensor)[0]
         probs = F.softmax(logits, dim=1)
         
-        # Classifier prediction
         top5_prob, top5_idx = torch.topk(probs, k=5)
         breeds = [dog_breeds[idx.item()] for idx in top5_idx[0]]
         probabilities = [prob.item() for prob in top5_prob[0]]
         
-        # Calculate relative probabilities
-        sum_probs = sum(probabilities[:3])  # 只取前三個來計算相對概率
+        sum_probs = sum(probabilities[:3])
         relative_probs = [f"{(prob/sum_probs * 100):.2f}%" for prob in probabilities[:3]]
         
-        # Debug output
-        print("\nClassifier Predictions:")
-        for breed, prob in zip(breeds[:5], probabilities[:5]):
-            print(f"{breed}: {prob:.4f}")
-            
         return probabilities[0], breeds[:3], relative_probs
         
         
 @spaces.GPU(duration=120)
 async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
     
-    if not hasattr(detect_multiple_dogs, 'model_yolo'):
-        detect_multiple_dogs.model_yolo = YOLO('yolov8l.pt')
-        
-    results = detect_multiple_dogs.model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
+    model_yolo = global_models['detection_model']
+    if model_yolo is None:
+        await initialize_models()
+        model_yolo = global_models['detection_model']
+    
+    results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
+    
     dogs = []
     boxes = []
     for box in results.boxes:
@@ -881,7 +901,7 @@ def show_details_html(choice, previous_output, initial_state):
 
 def main():
     with gr.Blocks(css=get_css_styles()) as iface:
-        # Header HTML
+        asyncio.create_task(initialize_models())
 
         gr.HTML("""
         <header style='text-align: center; padding: 20px; margin-bottom: 20px;'>