DurgaDeepak commited on
Commit
93d071e
·
verified ·
1 Parent(s): b207201

Update models/detection/detector.py

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