mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-02-05 18:55:23 +03:00
Fix issues in model generation script
This commit is contained in:
parent
f72a4e4d69
commit
3d7330a6a0
@ -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
|
||||||
|
|||||||
@ -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}/
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user