administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-02-05 10:45:21 +03:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Fix issues in model generation script

Browse Source
This commit is contained in:
Nate Meyer 2023-06-24 21:45:48 -04:00
parent 9ad2bd26a6
commit c9052f1ae4
2 changed files with 16 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
docker/support/tensorrt_detector/rootfs/etc/s6-overlay/s6-rc.d/trt-model-prepare/run
View File

@ -1,10 +1,10 @@
#!/command/with-contenv bash #!/command/with-contenv bash
# shellcheck shell=bash # shellcheck shell=bash
# Prepare the logs folder for s6-log # Generate models for the TensorRT detector
set -o errexit -o nounset -o pipefail set -o errexit -o nounset -o pipefail
OUTPUT_FOLDER?=/config/model_cache/tensorrt OUTPUT_FOLDER=${OUTPUT_FOLDER:-"/config/model_cache/tensorrt"}
# Create output folder # Create output folder
mkdir -p ${OUTPUT_FOLDER} mkdir -p ${OUTPUT_FOLDER}
@ -14,22 +14,22 @@ MODEL_CONVERT=""
for model in ${YOLO_MODELS//,/ } for model in ${YOLO_MODELS//,/ }
do do
if ![[ -f ${OUTPUT_FOLDER}/${model}.trt ]]; then if [[ ! -f ${OUTPUT_FOLDER}/${model}.trt ]]; then
if [[FIRST_MODEL == true]]; then if [[ ${FIRST_MODEL} = true ]]; then
MODEL_CONVERT="${model}" MODEL_CONVERT="${model}"
FIRST_MODEL=false FIRST_MODEL=false;
else else
MODEL_CONVERT+=",{$model}" MODEL_CONVERT+=",${model}";
fi fi
fi fi
done done
if [[${MODEL_CONVERT} == ""]]; then if [[ -z ${MODEL_CONVERT} ]]; then
echo "No models to convert." echo "No models to convert."
exit 0; exit 0
fi fi
echo "Generating the following TRT Models: ${YOLO_MODELS}" echo "Generating the following TRT Models: ${MODEL_CONVERT}"
# Build trt engine # Build trt engine
cd /usr/local/src/tensorrt_demos/yolo cd /usr/local/src/tensorrt_demos/yolo
@ -41,5 +41,5 @@ for model in ${MODEL_CONVERT//,/ }
do do
python3 yolo_to_onnx.py -m ${model} python3 yolo_to_onnx.py -m ${model}
python3 onnx_to_tensorrt.py -m ${model} python3 onnx_to_tensorrt.py -m ${model}
cp /tmp/tensorrt_demos-conditional_download/yolo/${model}.trt ${OUTPUT_FOLDER}/${model}.trt; cp ${model}.trt ${OUTPUT_FOLDER}/${model}.trt
done done

8
docker/support/tensorrt_detector/tensorrt_libyolo.sh
View File

@ -2,15 +2,17 @@
set -euxo pipefail set -euxo pipefail
SCRIPT_DIR="/usr/local/src/tensorrt_demos"
# Clone tensorrt_demos repo # Clone tensorrt_demos repo
git clone --depth 1 https://github.com/NateMeyer/tensorrt_demos.git -b conditional_download git clone --depth 1 https://github.com/NateMeyer/tensorrt_demos.git -b conditional_download
# Build libyolo # Build libyolo
cd ./tensorrt_demos/plugins && make all cd ./tensorrt_demos/plugins && make all
cp libyolo_layer.so /usr/local/lib/libyolo_layer.so cp libyolo_layer.so /usr/local/lib/libyolo_layer.so
cp libyolo_layer.so ../yolo/libyolo_layer.so
# Store yolo scripts for later conversion # Store yolo scripts for later conversion
cd ../ cd ../
mkdir -p /usr/local/src/tensorrt_demos mkdir -p ${SCRIPT_DIR}/plugins
cp -a yolo /usr/local/src/tensorrt_demos/ cp plugins/libyolo_layer.so ${SCRIPT_DIR}/plugins/libyolo_layer.so
cp -a yolo ${SCRIPT_DIR}/