DurgaDeepak commited on
Commit
d082050
·
verified ·
1 Parent(s): bbc95d9

Update models/detection/detector.py

Browse files
Files changed (1) hide show
  1. models/detection/detector.py +5 -9
models/detection/detector.py CHANGED
@@ -3,20 +3,16 @@ from PIL import Image, ImageDraw
3
  from huggingface_hub import hf_hub_download
4
  from ultralytics import YOLO
5
  import os
6
- import shutil
7
  import torch
8
 
9
  # Setup logger
10
  logger = logging.getLogger(__name__)
11
  logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
12
 
13
- # Optional: clear weights cache each time (only for dev use)
14
- shutil.rmtree("models/detection/weights", ignore_errors=True)
15
-
16
  class ObjectDetector:
17
  def __init__(self, model_key="yolov8n.pt", device="cpu"):
18
  """
19
- Initializes an Ultralytics YOLO model using HF download path.
20
 
21
  Args:
22
  model_key (str): e.g. 'yolov8n.pt', 'yolov8s.pt', etc.
@@ -50,16 +46,16 @@ class ObjectDetector:
50
  force_download=False
51
  )
52
 
53
- self.model = None # 🔁 Don't initialize on construction
54
- logger.info(f"Model path ready for {resolved_key}: {self.weights_path}")
55
 
56
  def load_model(self):
57
  if self.model is None:
58
- logger.info(" Loading YOLO model into memory...")
59
  self.model = YOLO(self.weights_path)
60
  if self.device == "cuda" and torch.cuda.is_available():
61
  self.model.to("cuda")
62
- logger.info(f"✅ YOLO model loaded on {self.device}")
63
 
64
  def predict(self, image: Image.Image, conf_threshold=0.25):
65
  self.load_model()
 
3
  from huggingface_hub import hf_hub_download
4
  from ultralytics import YOLO
5
  import os
 
6
  import torch
7
 
8
  # Setup logger
9
  logger = logging.getLogger(__name__)
10
  logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
11
 
 
 
 
12
  class ObjectDetector:
13
  def __init__(self, model_key="yolov8n.pt", device="cpu"):
14
  """
15
+ Initializes an Ultralytics YOLO model path, defers actual model loading.
16
 
17
  Args:
18
  model_key (str): e.g. 'yolov8n.pt', 'yolov8s.pt', etc.
 
46
  force_download=False
47
  )
48
 
49
+ logger.info(f"✅ YOLO weights ready for {resolved_key} at {self.weights_path}")
50
+ self.model = None # defer loading
51
 
52
  def load_model(self):
53
  if self.model is None:
54
+ logger.info("⚙️ Loading YOLO model into memory (runtime-safe)")
55
  self.model = YOLO(self.weights_path)
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()