mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-02-05 10:45:21 +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
|
||||
# shellcheck shell=bash
|
||||
# Prepare the logs folder for s6-log
|
||||
# Generate models for the TensorRT detector
|
||||
|
||||
set -o errexit -o nounset -o pipefail
|
||||
|
||||
OUTPUT_FOLDER?=/config/model_cache/tensorrt
|
||||
OUTPUT_FOLDER=${OUTPUT_FOLDER:-"/config/model_cache/tensorrt"}
|
||||
|
||||
# Create output folder
|
||||
mkdir -p ${OUTPUT_FOLDER}
|
||||
@ -14,22 +14,22 @@ MODEL_CONVERT=""
|
||||
|
||||
for model in ${YOLO_MODELS//,/ }
|
||||
do
|
||||
if ![[ -f ${OUTPUT_FOLDER}/${model}.trt ]]; then
|
||||
if [[FIRST_MODEL == true]]; then
|
||||
if [[ ! -f ${OUTPUT_FOLDER}/${model}.trt ]]; then
|
||||
if [[ ${FIRST_MODEL} = true ]]; then
|
||||
MODEL_CONVERT="${model}"
|
||||
FIRST_MODEL=false
|
||||
FIRST_MODEL=false;
|
||||
else
|
||||
MODEL_CONVERT+=",{$model}"
|
||||
MODEL_CONVERT+=",${model}";
|
||||
fi
|
||||
fi
|
||||
done
|
||||
|
||||
if [[${MODEL_CONVERT} == ""]]; then
|
||||
if [[ -z ${MODEL_CONVERT} ]]; then
|
||||
echo "No models to convert."
|
||||
exit 0;
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "Generating the following TRT Models: ${YOLO_MODELS}"
|
||||
echo "Generating the following TRT Models: ${MODEL_CONVERT}"
|
||||
|
||||
# Build trt engine
|
||||
cd /usr/local/src/tensorrt_demos/yolo
|
||||
@ -41,5 +41,5 @@ for model in ${MODEL_CONVERT//,/ }
|
||||
do
|
||||
python3 yolo_to_onnx.py -m ${model}
|
||||
python3 onnx_to_tensorrt.py -m ${model}
|
||||
cp /tmp/tensorrt_demos-conditional_download/yolo/${model}.trt ${OUTPUT_FOLDER}/${model}.trt;
|
||||
done
|
||||
cp ${model}.trt ${OUTPUT_FOLDER}/${model}.trt
|
||||
done
|
||||
|
||||
@ -2,15 +2,17 @@
|
||||
|
||||
set -euxo pipefail
|
||||
|
||||
SCRIPT_DIR="/usr/local/src/tensorrt_demos"
|
||||
|
||||
# Clone tensorrt_demos repo
|
||||
git clone --depth 1 https://github.com/NateMeyer/tensorrt_demos.git -b conditional_download
|
||||
|
||||
# Build libyolo
|
||||
cd ./tensorrt_demos/plugins && make all
|
||||
cp libyolo_layer.so /usr/local/lib/libyolo_layer.so
|
||||
cp libyolo_layer.so ../yolo/libyolo_layer.so
|
||||
|
||||
# Store yolo scripts for later conversion
|
||||
cd ../
|
||||
mkdir -p /usr/local/src/tensorrt_demos
|
||||
cp -a yolo /usr/local/src/tensorrt_demos/
|
||||
mkdir -p ${SCRIPT_DIR}/plugins
|
||||
cp plugins/libyolo_layer.so ${SCRIPT_DIR}/plugins/libyolo_layer.so
|
||||
cp -a yolo ${SCRIPT_DIR}/
|
||||
|
||||
Loading…
Reference in New Issue
Block a user