Update models/detection/detector.py

models/detection/detector.py

@@ -2,46 +2,61 @@ import logging
 from PIL import Image, ImageDraw
 from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
+import os
 import shutil
 
-
-
+# Setup logger
 logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+# Optional: clear weights cache each time (only for dev use)
 shutil.rmtree("models/detection/weights", ignore_errors=True)
+
 class ObjectDetector:
-    def __init__(self, model_key="yolov8n.pt", device="cpu"):
+    def __init__(self, model_key="yolov8n", device="cpu"):
         """
-        Initialize the Object Detection model using Ultralytics YOLO registry.
+        Initializes an Ultralytics YOLO model using HF download path.
 
         Args:
-            model_key (str): Model name supported by ultralytics, e.g. 'yolov5n', 'yolov8s', etc.
+            model_key (str): e.g. 'yolov8n', 'yolov8s', etc.
             device (str): 'cpu' or 'cuda'
         """
+        # Optional aliasing
         alias_map = {
+            "yolov8n": "yolov8n",
             "yolov8s": "yolov8s",
             "yolov8l": "yolov8l",
-            "yolov11b": "yolov11b",
+            "yolov11b": "yolov11b"
         }
 
-        raw_key = model_key.lower()
-        resolved_key = alias_map.get(raw_key, raw_key)
+        resolved_key = alias_map.get(model_key.lower(), model_key.lower())
 
-        self.device = device
-        self.model = YOLO(f"{resolved_key}.pt")
-        logger.info(f" Ultralytics YOLO model '{resolved_key}' initialized on {device}")
+        # HF repo map
+        hf_map = {
+            "yolov8n": ("ultralytics/yolov8", "yolov8n.pt"),
+            "yolov8s": ("ultralytics/yolov8", "yolov8s.pt"),
+            "yolov8l": ("ultralytics/yolov8", "yolov8l.pt"),
+            "yolov11b": ("Ultralytics/YOLO11", "yolov11b.pt"),
+        }
 
+        if resolved_key not in hf_map:
+            raise ValueError(f"Unsupported model key: {resolved_key}")
 
+        repo_id, filename = hf_map[resolved_key]
 
-    def predict(self, image: Image.Image, conf_threshold=0.25):
-        """
-        Run object detection.
+        # 📥 Download from HF Hub
+        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
+        )
 
-        Args:
-            image (PIL.Image.Image): Input image.
+        logger.info(f"✅ Loaded YOLO model: {resolved_key} from {weights_path}")
+        self.device = device
+        self.model = YOLO(weights_path)
 
-        Returns:
-            List[Dict]: List of detected objects with class name, confidence, and bbox.
-        """
+    def predict(self, image: Image.Image, conf_threshold=0.25):
         logger.info("Running object detection")
         results = self.model(image)
         detections = []
@@ -56,17 +71,6 @@ class ObjectDetector:
         return detections
 
     def draw(self, image: Image.Image, detections, alpha=0.5):
-        """
-        Draw bounding boxes on image.
-
-        Args:
-            image (PIL.Image.Image): Input image.
-            detections (List[Dict]): Detection results.
-            alpha (float): Blend strength.
-
-        Returns:
-            PIL.Image.Image: Image with bounding boxes drawn.
-        """
         overlay = image.copy()
         draw = ImageDraw.Draw(overlay)
         for det in detections: