Spaces:
Running
on
Zero
Running
on
Zero
Update models/detection/detector.py
Browse files
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="
|
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",
|