Fix issues in model generation script

This commit is contained in:
Nate Meyer 2023-06-24 21:45:48 -04:00
parent f72a4e4d69
commit 3d7330a6a0
2 changed files with 16 additions and 14 deletions

View File

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

View File

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