Update app.py

app.py

@@ -36,8 +36,6 @@ from ultralytics import YOLO
 from device_manager import DeviceManager, device_handler
 import asyncio
 import traceback
-import spaces
-
 
 # model_yolo = YOLO('yolov8l.pt')
 
@@ -622,8 +620,7 @@ def preprocess_image(image):
     ])
 
     return transform(image).unsqueeze(0)
-    
-@device_handler    
+       
 async def predict_single_dog(image):
     """
     Predicts the dog breed using only the classifier.
@@ -650,41 +647,49 @@ async def predict_single_dog(image):
             return probabilities[0], breeds[:3], relative_probs
             
     except RuntimeError as e:
-        if "out of memory" in str(e) or "CUDA" in str(e):
-            logger.warning("ZeroGPU unavailable, falling back to CPU")
+        if "out of memory" in str(e):
+            logger.warning("GPU memory exceeded, falling back to CPU")
             device_mgr._current_device = torch.device('cpu')
-            model.to('cpu')
-            return await predict_single_dog(image)  # 遞迴調用，使用 CPU
+            return await predict_single_dog(image)
         raise e
         
         
-@device_handler
 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
+    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
+        
+    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
+        
 
 def non_max_suppression(boxes, iou_threshold):
     keep = []