stuff

vr180_matting/detector.py

@@ -1,6 +1,4 @@
-import torch
 import numpy as np
-from ultralytics import YOLO
 from typing import List, Tuple, Dict, Any
 import cv2
 
@@ -13,14 +11,23 @@ class YOLODetector:
         self.confidence_threshold = confidence_threshold
         self.device = device
         self.model = None
-        self._load_model()
+        # Don't load model during init - load lazily when first used
     
     def _load_model(self):
-        """Load YOLOv8 model"""
+        """Load YOLOv8 model lazily"""
+        if self.model is not None:
+            return  # Already loaded
+            
         try:
+            # Import heavy dependencies only when needed
+            import torch
+            from ultralytics import YOLO
+            
             self.model = YOLO(f"{self.model_name}.pt")
             if self.device == "cuda" and torch.cuda.is_available():
                 self.model.to("cuda")
+                
+            print(f"🎯 Loaded YOLO model: {self.model_name}")
         except Exception as e:
             raise RuntimeError(f"Failed to load YOLO model {self.model_name}: {e}")
     
@@ -34,8 +41,9 @@ class YOLODetector:
         Returns:
             List of detection dictionaries with bbox, confidence, and class info
         """
+        # Load model lazily on first use
         if self.model is None:
-            raise RuntimeError("YOLO model not loaded")
+            self._load_model()
         
         results = self.model(frame, verbose=False)
         detections = []