Mirrors/blakeblackshear.frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2024-11-21 19:07:46 +01:00
Code Issues Projects Releases Wiki Activity

switch model download into bash script

Browse Source
This commit is contained in:
Indrek Mandre 2024-02-04 14:22:15 +02:00
parent 8a14703d26
commit 77fc38c84f
2 changed files with 27 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
docker/main/rootfs/etc/s6-overlay/s6-rc.d/download-models/run
View File

@ -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

62
frigate/util/download_models.py
View File

@ -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()