DurgaDeepak commited on
Commit
4cfdbcf
·
verified ·
1 Parent(s): dc71e71

Update models/detection/detector.py

Browse files
Files changed (1) hide show
  1. models/detection/detector.py +9 -0
models/detection/detector.py CHANGED
@@ -20,6 +20,15 @@ class ObjectDetector:
20
  weights_dir (str): Directory to store/download model weights.
21
  device (str): Inference device ("cpu" or "cuda").
22
  """
 
 
 
 
 
 
 
 
 
23
  model_key = model_key.lower() #to make model_key case-insensitive
24
  repo_map = {
25
  "yolov5n": ("ultralytics/yolov5", "yolov5n.pt"),
 
20
  weights_dir (str): Directory to store/download model weights.
21
  device (str): Inference device ("cpu" or "cuda").
22
  """
23
+ alias_map = {
24
+ "yolov5n-seg": "yolov5n",
25
+ "yolov5s-seg": "yolov5s",
26
+ "yolov8s": "yolov8s",
27
+ "yolov8l": "yolov8l",
28
+ "yolov11b": "yolov11b",
29
+ "rtdetr": "rtdetr"
30
+ }
31
+
32
  model_key = model_key.lower() #to make model_key case-insensitive
33
  repo_map = {
34
  "yolov5n": ("ultralytics/yolov5", "yolov5n.pt"),