Update models/detection/detector.py

models/detection/detector.py

@@ -3,6 +3,8 @@ import logging
 from PIL import Image, ImageDraw
 from huggingface_hub import hf_hub_download
 
+logger = logging.getLogger(__name__)
+
 class ObjectDetector:
     def __init__(self, model_key="yolov8n", device="cpu"):
         self.device = device
@@ -32,9 +34,9 @@ class ObjectDetector:
             import torch  # Safe to import here
             from ultralytics import YOLO  # Defer import
     
-            # FORCE: Avoid triggering CUDA by setting env var
-            os.environ["CUDA_VISIBLE_DEVICES"] = "-1" if self.device == "cpu" else ""
-    
+            if self.device == "cpu":
+                os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+                
             # Initialize model
             self.model = YOLO(self.weights_path)
     
@@ -58,3 +60,13 @@ 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)