Update models/detection/detector.py

models/detection/detector.py

@@ -4,6 +4,7 @@ from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
 import os
 import shutil
+import torch
 
 # Setup logger
 logger = logging.getLogger(__name__)
@@ -13,15 +14,16 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
 shutil.rmtree("models/detection/weights", ignore_errors=True)
 
 class ObjectDetector:
-    def __init__(self, model_key="yolov8n", device="cpu"):
+    def __init__(self, model_key="yolov8n.pt", device="cpu"):
         """
         Initializes an Ultralytics YOLO model using HF download path.
 
         Args:
-            model_key (str): e.g. 'yolov8n', 'yolov8s', etc.
+            model_key (str): e.g. 'yolov8n.pt', 'yolov8s.pt', etc.
             device (str): 'cpu' or 'cuda'
         """
-        # Optional aliasing
+        self.device = device
+        resolved_key = model_key.lower().replace(".pt", "")
         alias_map = {
             "yolov8n": "yolov8n",
             "yolov8s": "yolov8s",
@@ -29,9 +31,6 @@ class ObjectDetector:
             "yolov11b": "yolov11b"
         }
 
-        resolved_key = model_key.lower().replace(".pt", "")
-
-        # HF repo map
         hf_map = {
             "yolov8n": ("ultralytics/yolov8", "yolov8n.pt"),
             "yolov8s": ("ultralytics/yolov8", "yolov8s.pt"),
@@ -39,25 +38,32 @@ class ObjectDetector:
             "yolov11b": ("Ultralytics/YOLO11", "yolov11b.pt"),
         }
 
+        resolved_key = alias_map.get(resolved_key, resolved_key)
         if resolved_key not in hf_map:
             raise ValueError(f"Unsupported model key: {resolved_key}")
 
         repo_id, filename = hf_map[resolved_key]
-
-        # 📥 Download from HF Hub
-        weights_path = hf_hub_download(
+        self.weights_path = hf_hub_download(
             repo_id=repo_id,
             filename=filename,
             cache_dir="models/detection/weights",
-            force_download=True  # Optional: change to False for reuse
+            force_download=False
         )
 
-        logger.info(f"✅ Loaded YOLO model: {resolved_key} from {weights_path}")
-        self.device = device
-        self.model = YOLO(weights_path)
+        self.model = None  # 🔁 Don't initialize on construction
+        logger.info(f"Model path ready for {resolved_key}: {self.weights_path}")
+
+    def load_model(self):
+        if self.model is None:
+            logger.info("⏳ Loading YOLO model into memory...")
+            self.model = YOLO(self.weights_path)
+            if self.device == "cuda" and torch.cuda.is_available():
+                self.model.to("cuda")
+            logger.info(f"✅ YOLO model loaded on {self.device}")
 
     def predict(self, image: Image.Image, conf_threshold=0.25):
-        logger.info("Running object detection")
+        self.load_model()
+        logger.info("🔍 Running object detection")
         results = self.model(image)
         detections = []
         for r in results:
@@ -67,7 +73,7 @@ class ObjectDetector:
                     "confidence": float(box.conf),
                     "bbox": box.xyxy[0].tolist()
                 })
-        logger.info(f"Detected {len(detections)} objects")
+        logger.info(f"✅ Detected {len(detections)} objects")
         return detections
 
     def draw(self, image: Image.Image, detections, alpha=0.5):