Skip to content

Commit 842c39f

Browse files
authored
add yolov5n weight download util (#251)
* add yolov5n weight download util * format utils.yolov5
1 parent 50056d4 commit 842c39f

File tree

1 file changed

+17
-0
lines changed

1 file changed

+17
-0
lines changed

sahi/utils/yolov5.py

+17
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,30 @@
55

66

77
class Yolov5TestConstants:
8+
YOLOV5N_MODEL_URL = "https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5n.pt"
9+
YOLOV5N_MODEL_PATH = "tests/data/models/yolov5/yolov5n.pt"
10+
811
YOLOV5S6_MODEL_URL = "https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s6.pt"
912
YOLOV5S6_MODEL_PATH = "tests/data/models/yolov5/yolov5s6.pt"
1013

1114
YOLOV5M6_MODEL_URL = "https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5m6.pt"
1215
YOLOV5M6_MODEL_PATH = "tests/data/models/yolov5/yolov5m6.pt"
1316

1417

18+
def download_yolov5n_model(destination_path: Optional[str] = None):
19+
20+
if destination_path is None:
21+
destination_path = Yolov5TestConstants.YOLOV5N_MODEL_PATH
22+
23+
Path(destination_path).parent.mkdir(parents=True, exist_ok=True)
24+
25+
if not path.exists(destination_path):
26+
urllib.request.urlretrieve(
27+
Yolov5TestConstants.YOLOV5N_MODEL_URL,
28+
destination_path,
29+
)
30+
31+
1532
def download_yolov5s6_model(destination_path: Optional[str] = None):
1633

1734
if destination_path is None:

0 commit comments

Comments
 (0)