From 22bd9671b76c600d0e744a13d0a5a2db833319a6 Mon Sep 17 00:00:00 2001
From: Jan Gutsche <github@jaagut.de>
Date: Tue, 31 May 2022 16:54:12 +0200
Subject: [PATCH] Add download_weights.py

---
 README.md                   |  2 +-
 pyproject.toml              |  1 +
 scripts/download_weights.py | 33 +++++++++++++++++++++++++++++++++
 weights/download_weights.sh |  9 ---------
 4 files changed, 35 insertions(+), 10 deletions(-)
 create mode 100755 scripts/download_weights.py
 delete mode 100755 weights/download_weights.sh

diff --git a/README.md b/README.md
index 9f9c27c..5a13675 100644
--- a/README.md
+++ b/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
diff --git a/pyproject.toml b/pyproject.toml
index 7af144e..b39b9f7 100644
--- a/pyproject.toml
+++ b/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"
diff --git a/scripts/download_weights.py b/scripts/download_weights.py
new file mode 100755
index 0000000..72946ac
--- /dev/null
+++ b/scripts/download_weights.py
@@ -0,0 +1,33 @@
+#!/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.")
+            raise e
+
+
+if __name__ == '__main__':
+    run()
\ No newline at end of file
diff --git a/weights/download_weights.sh b/weights/download_weights.sh
deleted file mode 100755
index 8e0c664..0000000
--- a/weights/download_weights.sh
+++ /dev/null
@@ -1,9 +0,0 @@
-#!/bin/bash
-# Download weights for yoeo-rev-7
-wget "http://data.bit-bots.de/models/2021_12_06_flo_torso21_yoeo_7/yoeo.pth"
-# Download weights for vanilla YOLOv3
-wget -c "https://pjreddie.com/media/files/yolov3.weights" --header "Referer: pjreddie.com"
-# # Download weights for tiny YOLOv3
-wget -c "https://pjreddie.com/media/files/yolov3-tiny.weights" --header "Referer: pjreddie.com"
-# Download weights for backbone network
-wget -c "https://pjreddie.com/media/files/darknet53.conv.74" --header "Referer: pjreddie.com"