DurgaDeepak commited on
Commit
455843f
·
verified ·
1 Parent(s): d7e8ce4

Update models/detection/detector.py

Browse files
Files changed (1) hide show
  1. models/detection/detector.py +1 -0
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()