Update models/detection/detector.py

models/detection/detector.py

@@ -1,9 +1,7 @@
-import os
-import numpy as np
-from PIL import Image, ImageDraw
 import logging
+from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
-from utils.model_downloader import download_model_if_needed
+
 
 logger = logging.getLogger(__name__)
 
@@ -12,7 +10,7 @@ class ObjectDetector:
     Generalized Object Detection Wrapper for YOLOv5, YOLOv8, and future variants.
     """
 
-    def __init__(self, model_key="yolov5n-seg", device="cpu"):
+    def __init__(self, model_key="yolov5n", device="cpu"):
         """
         Initialize the Object Detection model.
 
@@ -21,10 +19,14 @@ class ObjectDetector:
             weights_dir (str): Directory to store/download model weights.
             device (str): Inference device ("cpu" or "cuda").
         """
+        model_key = model_key.lower() #to make model_key case-insensitive
         repo_map = {
             "yolov5n": ("ultralytics/yolov5", "yolov5n.pt"),
             "yolov8n": ("ultralytics/yolov8", "yolov8n.pt"),
-            # Add more if needed
+            "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: