Update models/detection/detector.py

models/detection/detector.py

@@ -9,18 +9,13 @@ import shutil
 logger = logging.getLogger(__name__)
 shutil.rmtree("models/detection/weights", ignore_errors=True)
 class ObjectDetector:
-    """
-    Generalized Object Detection Wrapper for YOLOv5, YOLOv8, and future variants.
-    """
-
     def __init__(self, model_key="yolov5n", device="cpu"):
         """
-        Initialize the Object Detection model.
+        Initialize the Object Detection model using Ultralytics YOLO registry.
 
         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").
+            model_key (str): Model name supported by ultralytics, e.g. 'yolov5n', 'yolov8s', etc.
+            device (str): 'cpu' or 'cuda'
         """
         alias_map = {
             "yolov5n-seg": "yolov5n",
@@ -28,39 +23,15 @@ class ObjectDetector:
             "yolov8s": "yolov8s",
             "yolov8l": "yolov8l",
             "yolov11b": "yolov11b",
-            "rtdetr": "rtdetr"
         }
 
         raw_key = model_key.lower()
-        model_key = alias_map.get(raw_key, raw_key) #to make model_key case-insensitive
-        repo_map = {
-            "yolov5n": ("ultralytics/yolov5", "yolov5n.pt"),
-            "yolov5s": ("ultralytics/yolov5", "yolov5s.pt"),
-            "yolov8n": ("ultralytics/yolov8", "yolov8n.pt"),
-            "yolov8s": ("ultralytics/yolov8", "yolov8s.pt"),
-            "yolov8l": ("ultralytics/yolov8", "yolov8l.pt"),
-            "yolov11b": ("Ultralytics/YOLO11", "yolov11b.pt"),
-            "rtdetr": ("IDEA-Research/RT-DETR", "rtdetr_r50vd_detr.pth")
-        }
-        
-        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",
-            force_download=True #Clear cache
-        )
-        
+        resolved_key = alias_map.get(raw_key, raw_key)
 
         self.device = device
-        print("Loading weights from:", weights_path)
-        self.model = YOLO(weights_path)
-        print("Model object type:", type(self.model))
-        print("Model class string:", self.model.__class__)
+        self.model = YOLO(resolved_key)
+        logger.info(f" Ultralytics YOLO model '{resolved_key}' initialized on {device}")
+
 
 
     def predict(self, image: Image.Image):