DurgaDeepak commited on
Commit
5caf904
·
verified ·
1 Parent(s): 8cfece3

Update models/detection/detector.py

Browse files
Files changed (1) hide show
  1. models/detection/detector.py +12 -2
models/detection/detector.py CHANGED
@@ -29,13 +29,23 @@ class ObjectDetector:
29
 
30
  def load_model(self):
31
  if self.model is None:
32
- from ultralytics import YOLO # Deferring import to avoid early CUDA trigger
 
 
 
 
 
 
33
  self.model = YOLO(self.weights_path)
 
 
34
  if self.device == "cuda" and torch.cuda.is_available():
35
- self.model.to("cuda")
 
36
  return self
37
 
38
 
 
39
  def predict(self, image: Image.Image, conf_threshold=0.25):
40
  self.load_model()
41
  results = self.model(image)
 
29
 
30
  def load_model(self):
31
  if self.model is None:
32
+ import torch # Safe to import here
33
+ from ultralytics import YOLO # Defer import
34
+
35
+ # FORCE: Avoid triggering CUDA by setting env var
36
+ os.environ["CUDA_VISIBLE_DEVICES"] = "-1" if self.device == "cpu" else ""
37
+
38
+ # Initialize model
39
  self.model = YOLO(self.weights_path)
40
+
41
+ # Move to CUDA only if necessary and safe
42
  if self.device == "cuda" and torch.cuda.is_available():
43
+ self.model.to("cuda")
44
+
45
  return self
46
 
47
 
48
+
49
  def predict(self, image: Image.Image, conf_threshold=0.25):
50
  self.load_model()
51
  results = self.model(image)