Update models/detection/detector.py

models/detection/detector.py

@@ -1,73 +1,86 @@
-import os
-import numpy as np
-from PIL import Image, ImageDraw
-import logging
-from ultralytics import YOLO
-from utils.model_downloader import download_model_if_needed
-
-logger = logging.getLogger(__name__)
-
-class ObjectDetector:
-    """
-    Generalized Object Detection Wrapper for YOLOv5, YOLOv8, and future variants.
-    """
-
-    def __init__(self, model_key="yolov5n-seg", weights_dir="models/detection/weights", device="cpu"):
-        """
-        Initialize the Object Detection model.
-
-        Args:
-            model_key (str): Model identifier as defined in model_downloader.py.
-            weights_dir (str): Directory to store/download model weights.
-            device (str): Inference device ("cpu" or "cuda").
-        """
-        weights_path = os.path.join(weights_dir, f"{model_key}.pt")
-        download_model_if_needed(model_key, weights_path)
-
-        logger.info(f"Loading Object Detection model '{model_key}' from {weights_path}")
-        self.device = device
-        self.model = YOLO(weights_path)
-
-    def predict(self, image: Image.Image):
-        """
-        Run object detection.
-
-        Args:
-            image (PIL.Image.Image): Input image.
-
-        Returns:
-            List[Dict]: List of detected objects with class name, confidence, and bbox.
-        """
-        logger.info("Running object detection")
-        results = self.model(image)
-        detections = []
-        for r in results:
-            for box in r.boxes:
-                detections.append({
-                    "class_name": r.names[int(box.cls)],
-                    "confidence": float(box.conf),
-                    "bbox": box.xyxy[0].tolist()
-                })
-        logger.info(f"Detected {len(detections)} objects")
-        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:
-            bbox = det["bbox"]
-            label = f'{det["class_name"]} {det["confidence"]:.2f}'
-            draw.rectangle(bbox, outline="red", width=2)
-            draw.text((bbox[0], bbox[1]), label, fill="red")
-        return Image.blend(image, overlay, alpha)
+import os
+import numpy as np
+from PIL import Image, ImageDraw
+import logging
+from ultralytics import YOLO
+from utils.model_downloader import download_model_if_needed
+
+logger = logging.getLogger(__name__)
+
+class ObjectDetector:
+    """
+    Generalized Object Detection Wrapper for YOLOv5, YOLOv8, and future variants.
+    """
+
+    def __init__(self, model_key="yolov5n-seg", device="cpu"):
+        """
+        Initialize the Object Detection model.
+
+        Args:
+            model_key (str): Model identifier as defined in model_downloader.py.
+            weights_dir (str): Directory to store/download model weights.
+            device (str): Inference device ("cpu" or "cuda").
+        """
+        repo_map = {
+            "yolov5n": ("ultralytics/yolov5", "yolov5n.pt"),
+            "yolov8n": ("ultralytics/yolov8", "yolov8n.pt"),
+            # Add more if needed
+        }
+        
+        if model_key not in repo_map:
+            raise ValueError(f"Unsupported model_key: {model_key}")
+
+        repo_id, filename = repo_map[model_key]
+
+        weights_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            cache_dir="models/detection/weights"
+        )
+
+        self.device = device
+        self.model = YOLO(weights_path)
+
+    def predict(self, image: Image.Image):
+        """
+        Run object detection.
+
+        Args:
+            image (PIL.Image.Image): Input image.
+
+        Returns:
+            List[Dict]: List of detected objects with class name, confidence, and bbox.
+        """
+        logger.info("Running object detection")
+        results = self.model(image)
+        detections = []
+        for r in results:
+            for box in r.boxes:
+                detections.append({
+                    "class_name": r.names[int(box.cls)],
+                    "confidence": float(box.conf),
+                    "bbox": box.xyxy[0].tolist()
+                })
+        logger.info(f"Detected {len(detections)} objects")
+        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:
+            bbox = det["bbox"]
+            label = f'{det["class_name"]} {det["confidence"]:.2f}'
+            draw.rectangle(bbox, outline="red", width=2)
+            draw.text((bbox[0], bbox[1]), label, fill="red")
+        return Image.blend(image, overlay, alpha)