Update app.py

app.py

@@ -503,9 +503,6 @@ import traceback
 #     iface.launch()
 
 
-model_yolo = YOLO('yolov8l.pt')
-
-history_manager = UserHistoryManager()
 
 dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staffordshire_Terrier",
               "Appenzeller", "Australian_Terrier", "Bedlington_Terrier", "Bernese_Mountain_Dog", "Bichon_Frise",
@@ -537,6 +534,8 @@ dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staff
 
 device_mgr = DeviceManager()
 
+history_manager = UserHistoryManager()
+
 class MultiHeadAttention(nn.Module):
 
     def __init__(self, in_dim, num_heads=8):
@@ -597,15 +596,18 @@ num_classes = len(dog_breeds)
 
 # Initialize base model
 model = BaseModel(num_classes=num_classes)
-
+model = device_mgr.to_device(model)
 # Load model path
 model_path = '124_best_model_dog.pth'
-checkpoint = torch.load(model_path, map_location=device_mgr.get_optimal_device())
+checkpoint = torch.load(model_path, map_location=device_mgr.get_device(), weights_only=True)
 
 # Load model state
 model.load_state_dict(checkpoint['base_model'], strict=False)
 model.eval()
 
+model_yolo = YOLO('yolov8l.pt')
+model_yolo = device_mgr.to_device(model_yolo)
+
 # Image preprocessing function
 def preprocess_image(image):
     # If the image is numpy.ndarray turn into PIL.Image
@@ -621,74 +623,59 @@ def preprocess_image(image):
 
     return transform(image).unsqueeze(0)
        
+
+@adaptive_gpu(duration=30)
 async def predict_single_dog(image):
-    """
-    Predicts the dog breed using only the classifier.
-    """
-    try:
-        image_tensor = preprocess_image(image).to(device_mgr.get_optimal_device())
+    """單獨的狗預測函數"""
+    image_tensor = preprocess_image(image)
+    image_tensor = device_mgr.to_device(image_tensor)
+    
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        logits = outputs[0] if isinstance(outputs, tuple) else outputs
+        probs = F.softmax(logits, dim=1)
         
-        with torch.no_grad():
-            outputs = model(image_tensor)  # 同步調用
-            logits = outputs[0] if isinstance(outputs, tuple) else outputs
-            probs = F.softmax(logits, dim=1)
-            
-            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]]
-            
-            sum_probs = sum(probabilities[:3])
-            relative_probs = [f"{(prob/sum_probs * 100):.2f}%" for prob in probabilities[:3]]
-            
-            print("\nClassifier Predictions:")
-            for breed, prob in zip(breeds[:5], probabilities[:5]):
-                print(f"{breed}: {prob:.4f}")
-                
-            return probabilities[0], breeds[:3], relative_probs
-            
-    except RuntimeError as e:
-        if "out of memory" in str(e):
-            logger.warning("GPU memory exceeded, falling back to CPU")
-            device_mgr._current_device = torch.device('cpu')
-            return await predict_single_dog(image)
-        raise e
+        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]]
         
+        sum_probs = sum(probabilities[:3])
+        relative_probs = [f"{(prob/sum_probs * 100):.2f}%" for prob in probabilities[:3]]
         
-async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
-    try:     
-        results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
-        dogs = []
-        boxes = []
-        for box in results.boxes:
-            if box.cls == 16:  # COCO dataset class for dog is 16
-                xyxy = box.xyxy[0].tolist()
-                confidence = box.conf.item()
-                boxes.append((xyxy, confidence))
-    
-        if not boxes:
-            dogs.append((image, 1.0, [0, 0, image.width, image.height]))
-        else:
-            nms_boxes = non_max_suppression(boxes, iou_threshold)
-    
-            for box, confidence in nms_boxes:
-                x1, y1, x2, y2 = box
-                w, h = x2 - x1, y2 - y1
-                x1 = max(0, x1 - w * 0.05)
-                y1 = max(0, y1 - h * 0.05)
-                x2 = min(image.width, x2 + w * 0.05)
-                y2 = min(image.height, y2 + h * 0.05)
-                cropped_image = image.crop((x1, y1, x2, y2))
-                dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
-    
-        return dogs
+        print("\nClassifier Predictions:")
+        for breed, prob in zip(breeds[:5], probabilities[:5]):
+            print(f"{breed}: {prob:.4f}")
+            
+        return probabilities[0], breeds[:3], relative_probs
         
-    except RuntimeError as e:
-        if "out of memory" in str(e):
-            logger.warning("GPU memory exceeded, falling back to CPU")
-            device_mgr._current_device = torch.device('cpu')
-            # 重新嘗試檢測
-            return await detect_multiple_dogs(image, conf_threshold, iou_threshold)
-        raise e
+@adaptive_gpu(duration=30)      
+async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
+    """複數狗預測函數"""
+    results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
+    dogs = []
+    boxes = []
+    for box in results.boxes:
+        if box.cls == 16:  # COCO dataset class for dog is 16
+            xyxy = box.xyxy[0].tolist()
+            confidence = box.conf.item()
+            boxes.append((xyxy, confidence))
+
+    if not boxes:
+        dogs.append((image, 1.0, [0, 0, image.width, image.height]))
+    else:
+        nms_boxes = non_max_suppression(boxes, iou_threshold)
+
+        for box, confidence in nms_boxes:
+            x1, y1, x2, y2 = box
+            w, h = x2 - x1, y2 - y1
+            x1 = max(0, x1 - w * 0.05)
+            y1 = max(0, y1 - h * 0.05)
+            x2 = min(image.width, x2 + w * 0.05)
+            y2 = min(image.height, y2 + h * 0.05)
+            cropped_image = image.crop((x1, y1, x2, y2))
+            dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
+
+    return dogs
         
 
 def non_max_suppression(boxes, iou_threshold):