#!/usr/bin/env sh

set -e

: "${OVERRIDE:=false}"

: "${DEBUG:=false}"
[ "${DEBUG}" = "true" ] && set -x

: "${MODEL_TYPE:=yolov5s}"
: "${UPDATE_YOLOV5}"

: "${GPU:=false}"

BASEDIR="$(cd "$(dirname "$0")" && pwd)"

[ "${UPDATE_YOLOV5}" = "true" ] && [ -d "$BASEDIR/yolov5" ] && \rm -rf yolov5

main() {
    prev_dir="$(pwd)"
    cd "$BASEDIR"

    ! [ -d yolov5 ] &&
    mkdir -p yolov5 && cd yolov5 &&
    ! [ -d yolov5.git ] && git clone --bare https://github.com/ultralytics/yolov5.git &&
    git --git-dir=./yolov5.git --work-tree=. checkout HEAD -- data models utils LICENSE detect.py export.py val.py requirements.txt && \rm -rf yolov5.git && cd "${BASEDIR}"

    exists=0

    [ -d "${BASEDIR}/${MODEL_TYPE}_saved_model" ] && exists=1
    [ "${OVERRIDE}" = "true" ] && \rm -rf "${BASEDIR}/${MODEL_TYPE}_saved_model"
    ! [ -d "${BASEDIR}/model_repository/tensorflow_${MODEL_TYPE}/1" ] && mkdir -p "${BASEDIR}/model_repository/tensorflow_${MODEL_TYPE}/1"
    [ -d "${BASEDIR}/model_repository/tensorflow_${MODEL_TYPE}/1/model.savedmodel" ] && \rm -rf "$BASEDIR/model_repository/tensorflow_${MODEL_TYPE}/1/model.savedmodel"

    if [ "${exists}" -eq 0 ] || [ "$OVERRIDE" = "true" ]; then
		if [ "$GPU" = "true" ]; then
			python3 yolov5/export.py --weights "$MODEL_TYPE.pt" --include saved_model torchscript --batch-size "$1" --dynamic --device 0
		else
			python3 yolov5/export.py --weights "$MODEL_TYPE.pt" --include saved_model torchscript --batch-size "$1" --dynamic
		fi
    fi
    cp -r "${BASEDIR}/${MODEL_TYPE}_saved_model" "${BASEDIR}/model_repository/tensorflow_${MODEL_TYPE}/1/model.savedmodel"

    echo "Successfully export YOLOv5 weights for SavedModel."
    cd "${prev_dir}"
}

if ! [ "$#" -eq 1 ]; then
    echo "Usage: $0 <batch_size>. Set to 1 by default if not passed."
    main 1
else
    main "$@"
fi
