Update models/detection/detector.py

models/detection/detector.py

@@ -1,31 +1,34 @@
-import logging
-from PIL import Image, ImageDraw
-from huggingface_hub import hf_hub_download
-from ultralytics import YOLO
-import os
-
-logger = logging.getLogger(__name__)
-
 class ObjectDetector:
-    def __init__(self, model_key="yolov8n.pt", device="cpu"):
+    def __init__(self, model_key="yolov8n", device="cpu"):
         self.device = device
         self.model = None
         self.model_key = model_key.lower().replace(".pt", "")
-        self.repo_map = {
+
+        hf_map = {
             "yolov8n": ("ultralytics/yolov8", "yolov8n.pt"),
             "yolov8s": ("ultralytics/yolov8", "yolov8s.pt"),
             "yolov8l": ("ultralytics/yolov8", "yolov8l.pt"),
             "yolov11b": ("Ultralytics/YOLO11", "yolov11b.pt"),
         }
 
-    def load_model(self):
-        if self.model is not None:
-            return
-        if self.model_key not in self.repo_map:
+        if self.model_key not in hf_map:
             raise ValueError(f"Unsupported model key: {self.model_key}")
-        repo_id, filename = self.repo_map[self.model_key]
-        weights_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir="models/detection/weights")
-        self.model = YOLO(weights_path)  # ✅ ZeroGPU-safe: runtime only
+
+        repo_id, filename = hf_map[self.model_key]
+        self.weights_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            cache_dir="models/detection/weights",
+            force_download=False
+        )
+
+    def load_model(self):
+        if self.model is None:
+            from ultralytics import YOLO  # Defer import
+            self.model = YOLO(self.weights_path)
+            if self.device == "cuda":
+                self.model.to("cuda")
+        return self  # So you can chain
 
     def predict(self, image: Image.Image, conf_threshold=0.25):
         self.load_model()
@@ -39,13 +42,3 @@ class ObjectDetector:
                     "bbox": box.xyxy[0].tolist()
                 })
         return detections
-
-    def draw(self, image: Image.Image, detections, alpha=0.5):
-        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)