Add download_weights.py

README.md

@@ -28,7 +28,7 @@ Also have a look at the other installing method, if you want to use the commands
 #### Download pretrained weights
 
 ```bash
-./weights/download_weights.sh
+poetry run yoeo-download-weights
 ```
 
 ## Test

pyproject.toml

@@ -34,3 +34,4 @@ build-backend = "poetry.core.masonry.api"
 yoeo-detect = "yoeo.detect:run"
 yoeo-train = "yoeo.train:run"
 yoeo-test = "yoeo.test:run"
+yoeo-download-weights = "scripts.download_weights:run"

scripts/download_weights.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+import os
+import pathlib
+import argparse
+import urllib.request
+
+
+def run():
+ parser = argparse.ArgumentParser(description='Download YOEO pretrained weights')
+ parser.add_argument(
+ '--output',
+ '-o',
+ type=pathlib.Path,
+ default='weights/yoeo.pth',
+ help='The pretrained weights file (.pth) will be written to this path. Defaults to "weights/yoeo.pth"',
+ )
+ args = parser.parse_args()
+
+ url = "https://data.bit-bots.de/models/2021_12_06_flo_torso21_yoeo_7/yoeo.pth"
+ output_path = args.output
+
+ with urllib.request.urlopen(url) as input_file:
+ try:
+ with open(output_path, 'xb') as output_file:
+ print(f"Saving pretrained weights to: {output_path}")
+ output_file.write(input_file.read())
+ except FileExistsError as e:
+ print(f"ERROR: The output file {output_path} does already exist. Will abort and not overwrite.")
+
+
+if __name__ == '__main__':
+ run()