Spaces:
Running
on
Zero
Running
on
Zero
Update models/detection/detector.py
Browse files
models/detection/detector.py
CHANGED
@@ -30,6 +30,7 @@ class ObjectDetector:
|
|
30 |
)
|
31 |
|
32 |
def load_model(self):
|
|
|
33 |
if self.model is None:
|
34 |
import torch # Safe to import here
|
35 |
from ultralytics import YOLO # Defer import
|
@@ -50,6 +51,8 @@ class ObjectDetector:
|
|
50 |
|
51 |
def predict(self, image: Image.Image, conf_threshold=0.25):
|
52 |
self.load_model()
|
|
|
|
|
53 |
results = self.model(image)
|
54 |
detections = []
|
55 |
for r in results:
|
|
|
30 |
)
|
31 |
|
32 |
def load_model(self):
|
33 |
+
logger.info(f"Loading model from path: {self.weights_path}")
|
34 |
if self.model is None:
|
35 |
import torch # Safe to import here
|
36 |
from ultralytics import YOLO # Defer import
|
|
|
51 |
|
52 |
def predict(self, image: Image.Image, conf_threshold=0.25):
|
53 |
self.load_model()
|
54 |
+
if self.model is None:
|
55 |
+
raise RuntimeError("YOLO model not loaded. Call load_model() first.")
|
56 |
results = self.model(image)
|
57 |
detections = []
|
58 |
for r in results:
|