add yolov5n weight download util (#251)

* add yolov5n weight download util * format utils.yolov5
obss · Oct 19, 2021 · 842c39f · 842c39f
1 parent 50056d4
commit 842c39f
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/sahi/utils/yolov5.py b/sahi/utils/yolov5.py
@@ -5,13 +5,30 @@
 
 
 class Yolov5TestConstants:
+    YOLOV5N_MODEL_URL = "https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5n.pt"
+    YOLOV5N_MODEL_PATH = "tests/data/models/yolov5/yolov5n.pt"
+
     YOLOV5S6_MODEL_URL = "https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s6.pt"
     YOLOV5S6_MODEL_PATH = "tests/data/models/yolov5/yolov5s6.pt"
 
     YOLOV5M6_MODEL_URL = "https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5m6.pt"
     YOLOV5M6_MODEL_PATH = "tests/data/models/yolov5/yolov5m6.pt"
 
 
+def download_yolov5n_model(destination_path: Optional[str] = None):
+
+    if destination_path is None:
+        destination_path = Yolov5TestConstants.YOLOV5N_MODEL_PATH
+
+    Path(destination_path).parent.mkdir(parents=True, exist_ok=True)
+
+    if not path.exists(destination_path):
+        urllib.request.urlretrieve(
+            Yolov5TestConstants.YOLOV5N_MODEL_URL,
+            destination_path,
+        )
+
+
 def download_yolov5s6_model(destination_path: Optional[str] = None):
 
     if destination_path is None: