DurgaDeepak commited on
Commit
9c37b92
·
verified ·
1 Parent(s): fddc40c

Update models/detection/detector.py

Browse files
Files changed (1) hide show
  1. models/detection/detector.py +1 -3
models/detection/detector.py CHANGED
@@ -9,7 +9,7 @@ import shutil
9
  logger = logging.getLogger(__name__)
10
  shutil.rmtree("models/detection/weights", ignore_errors=True)
11
  class ObjectDetector:
12
- def __init__(self, model_key="yolov5n", device="cpu"):
13
  """
14
  Initialize the Object Detection model using Ultralytics YOLO registry.
15
 
@@ -18,8 +18,6 @@ class ObjectDetector:
18
  device (str): 'cpu' or 'cuda'
19
  """
20
  alias_map = {
21
- "yolov5n-seg": "yolov5n",
22
- "yolov5s-seg": "yolov5s",
23
  "yolov8s": "yolov8s",
24
  "yolov8l": "yolov8l",
25
  "yolov11b": "yolov11b",
 
9
  logger = logging.getLogger(__name__)
10
  shutil.rmtree("models/detection/weights", ignore_errors=True)
11
  class ObjectDetector:
12
+ def __init__(self, model_key="yolov8n", device="cpu"):
13
  """
14
  Initialize the Object Detection model using Ultralytics YOLO registry.
15
 
 
18
  device (str): 'cpu' or 'cuda'
19
  """
20
  alias_map = {
 
 
21
  "yolov8s": "yolov8s",
22
  "yolov8l": "yolov8l",
23
  "yolov11b": "yolov11b",