#!/command/with-contenv bash
# shellcheck shell=bash
# Generate models for the TensorRT detector

set -o errexit -o nounset -o pipefail

MODEL_CACHE_DIR=${MODEL_CACHE_DIR:-"/config/model_cache/tensorrt"}
OUTPUT_FOLDER="${MODEL_CACHE_DIR}/${TRT_VER}"

# Create output folder
mkdir -p ${OUTPUT_FOLDER}

FIRST_MODEL=true
MODEL_CONVERT=""

for model in ${YOLO_MODELS//,/ }
do
    # Remove old link in case path/version changed
    rm -f ${MODEL_CACHE_DIR}/${model}.trt
    
    if [[ ! -f ${OUTPUT_FOLDER}/${model}.trt ]]; then
        if [[ ${FIRST_MODEL} = true ]]; then
            MODEL_CONVERT="${model}"
            FIRST_MODEL=false;
        else
            MODEL_CONVERT+=",${model}";
        fi
    else
        ln -s ${OUTPUT_FOLDER}/${model}.trt ${MODEL_CACHE_DIR}/${model}.trt
    fi
done

if [[ -z ${MODEL_CONVERT} ]]; then
    echo "No models to convert."
    exit 0
fi

echo "Generating the following TRT Models: ${MODEL_CONVERT}"

# Build trt engine
cd /usr/local/src/tensorrt_demos/yolo

# Download yolo weights
./download_yolo.sh $MODEL_CONVERT > /dev/null

for model in ${MODEL_CONVERT//,/ }
do
    echo "Converting ${model} model"
    python3 yolo_to_onnx.py -m ${model} > /dev/null
    python3 onnx_to_tensorrt.py -m ${model} > /dev/null
    cp ${model}.trt ${OUTPUT_FOLDER}/${model}.trt
    ln -s ${OUTPUT_FOLDER}/${model}.trt ${MODEL_CACHE_DIR}/${model}.trt
done
