From 77fc38c84f16325b301e6822652e2411642c5dad Mon Sep 17 00:00:00 2001
From: Indrek Mandre <indrek@mare.ee>
Date: Sun, 4 Feb 2024 14:22:15 +0200
Subject: [PATCH] switch model download into bash script

---
 .../s6-overlay/s6-rc.d/download-models/run    | 28 ++++++++-
 frigate/util/download_models.py               | 62 -------------------
 2 files changed, 27 insertions(+), 63 deletions(-)
 delete mode 100644 frigate/util/download_models.py

diff --git a/docker/main/rootfs/etc/s6-overlay/s6-rc.d/download-models/run b/docker/main/rootfs/etc/s6-overlay/s6-rc.d/download-models/run
index 4a7abe814..7f28e30d0 100755
--- a/docker/main/rootfs/etc/s6-overlay/s6-rc.d/download-models/run
+++ b/docker/main/rootfs/etc/s6-overlay/s6-rc.d/download-models/run
@@ -4,5 +4,31 @@
 
 set -o errexit -o nounset -o pipefail
 
-python3 /opt/frigate/frigate/util/download_models.py
+MODEL_CACHE_DIR=${MODEL_CACHE_DIR:-"/config/model_cache"}
+YOLOV8_DIR="$MODEL_CACHE_DIR/yolov8"
+YOLOV8_URL=https://github.com/harakas/models/releases/download/yolov8.1-1.0/yolov8.small.models.tar.gz
+YOLOV8_DIGEST=3407d38e44163e84f197827d5fc61b31f0e3dd82de90a15838e7faf9a18f6278
+
+if [ "$DOWNLOAD_YOLOV8" = "1" ]; then
+  echo "download-models: DOWNLOAD_YOLOV8=${DOWNLOAD_YOLOV8}, running download"
+  if ! test -f "${YOLOV8_DIR}/model.fetched"; then
+    mkdir -p $YOLOV8_DIR
+    TMP_FILE="${YOLOV8_DIR}/download.tar.gz"
+    curl --no-progress-meter -L --max-filesize 500M --insecure --output $TMP_FILE "${YOLOV8_URL}"
+    digest=$(sha256sum $TMP_FILE | awk '{print $1}')
+    if [ "$digest" = "$YOLOV8_DIGEST" ]; then
+      echo "download-models: Extracting downloaded file"
+      cd $YOLOV8_DIR
+      tar zxf $TMP_FILE
+      rm $TMP_FILE
+      touch model.fetched
+      echo "download-models: Yolov8 download done, files placed into ${YOLOV8_DIR}"
+    else
+      echo "download-models: Downloaded file digest does not match: got $digest, expected $YOLOV8_DIGEST"
+      rm $TMP_FILE
+    fi
+  else
+    echo "download-models: ${YOLOV8_DIR}/model.fetched already present"
+  fi
+fi
 
diff --git a/frigate/util/download_models.py b/frigate/util/download_models.py
deleted file mode 100644
index c0740bcd9..000000000
--- a/frigate/util/download_models.py
+++ /dev/null
@@ -1,62 +0,0 @@
-
-import os
-import hashlib
-import logging
-
-logger = logging.getLogger(__name__)
-
-def execute(cmd):
-    logger.info(f'executing: {cmd}')
-    ret = os.system(cmd)
-    if ret != 0:
-        logger.error(f'command {cmd} returned exit code: {ret}')
-    return ret
-
-def file_digest(filename):
-    with open(filename, mode='rb') as f:
-        d = hashlib.sha256()
-        while True:
-            buf = f.read(131072)
-            if not buf:
-                break
-            d.update(buf)
-        return d.hexdigest()
-
-def download_yolov8():
-    url = 'https://github.com/harakas/models/releases/download/yolov8.1-1.0/yolov8.small.models.tar.gz'
-    target_dir = '/config/model_cache/yolov8'
-    download_name = 'download.tar.gz'
-    allowed_hashes = set(('3407d38e44163e84f197827d5fc61b31f0e3dd82de90a15838e7faf9a18f6278',))
-    if os.path.exists(f'{target_dir}/model.fetched'):
-        logger.info(f"{target_dir}/model.fetched present, nothing to do")
-        return
-    os.makedirs(target_dir, exist_ok=True)
-    os.chdir(target_dir)
-    ret = execute(f'curl --silent -L --max-filesize 500M --insecure --output {target_dir}/{download_name} {url}')
-    if ret != 0:
-        try:
-            os.unlink(download_name)
-        except:
-            pass
-        return
-    digest = file_digest(download_name)
-    if digest not in allowed_hashes:
-        logger.info(f"Downloaded yolov8 model file not in allowed hashes: {digest}")
-        os.unlink(download_name)
-        return
-    ret = execute(f'tar zxf {download_name}')
-    if ret != 0:
-        os.unlink(download_name)
-        return
-    logger.info(f"Download and and extraction complete, marking download as done")
-    open('model.fetched', 'w').close()
-    os.unlink(download_name)
-
-
-if __name__ == '__main__':
-    if os.getenv('DOWNLOAD_YOLOV8') in ('1', 'TRUE'):
-        logging.basicConfig(level=logging.INFO)
-        logger = logging.getLogger(__file__)
-        logger.info(f"DOWNLOAD_YOLOV8={os.getenv('DOWNLOAD_YOLOV8')} -- running yolov8 download")
-        download_yolov8()
-