Spaces:
Running
on
Zero
Running
on
Zero
Update models/detection/detector.py
Browse files
models/detection/detector.py
CHANGED
@@ -56,6 +56,7 @@ class ObjectDetector:
|
|
56 |
if self.device == "cuda" and torch.cuda.is_available():
|
57 |
self.model.to("cuda")
|
58 |
logger.info(f"✅ YOLO model initialized on {self.device}")
|
|
|
59 |
|
60 |
def predict(self, image: Image.Image, conf_threshold=0.25):
|
61 |
self.load_model()
|
|
|
56 |
if self.device == "cuda" and torch.cuda.is_available():
|
57 |
self.model.to("cuda")
|
58 |
logger.info(f"✅ YOLO model initialized on {self.device}")
|
59 |
+
return self
|
60 |
|
61 |
def predict(self, image: Image.Image, conf_threshold=0.25):
|
62 |
self.load_model()
|