yolo_web/scripts/train.sh

set -e

if [ -f .env ]; then
	source .env
fi

if [ -z ${1} ]; then
	echo "Please provide a JSON object"
	exit 1
fi

function fail () {
	ID="${1}"
	META=${2}
	URL="${YOLO_WEB_URL}/job/fail/${ID}"
	echo "${META}"
	curl -s -X POST \
		-H 'Content-Type: application/json' \
		-d "{\"meta\":\"${META}\"}" "${URL}"
}

trap 'fail $ID' ERR

JSON="${1}"

ID=$(echo $JSON | jq -r '.id')
DATASET=$(echo $JSON | jq -r '.dataset')
MODEL=$(echo $JSON | jq -r '.model')
NAME=$(echo $JSON | jq -r '.name')
FILEPATH=$(echo $JSON | jq -r '.path')

DOWNLOAD="${YOLO_WEB_URL}${FILEPATH}"
DEST="${DATASETS}/${DATASET}.zip"
UNZIPPED="${DATASETS}/${DATASET}"

echo "Downloading ${DOWNLOAD}"

if [ ! -f "${DEST}" ]; then
	wget -q -O "${DEST}" "${DOWNLOAD}"
fi

unzip -o "${DEST}" -d "${UNZIPPED}"

if [ ! -f "${UNZIPPED}/data.yaml" ]; then
	fail "${ID}" "Invalid dataset"
	rm -rf "${UNZIPPED}"
	rm "${DEST}"
	exit 1
fi

FULLDEST=$(realpath "${UNZIPPED}")
python scripts/data.py "${FULLDEST}/data.yaml"

cd yolov5
source env/bin/activate

WEIGHTS="yolov5s.pt"

if [[ "${MODEL}" == "yolov5m_onnx" ]]; then
	WEIGHTS="yolov5m.pt"
elif [[ "${MODEL}" == "yolov5l_onnx" ]]; then
	WEIGHTS="yolov5l.pt"
fi

python train.py --img 640 --batch 16 --epochs 300 --data "${FULLDEST}/data.yaml" --weights "${WEIGHTS}" --cache --project ../models/ --name "${ID}"

if [ ! -f "../models/${ID}/weights/best.pt" ]; then
	fail "${ID}" "Error training model"
	rm -rf "${UNZIPPED}"
	rm "${DEST}"
	exit 3
fi

python export.py --weights "../models/${ID}/weights/best.pt" --include onnx

if [ ! -f "../models/${ID}/weights/best.onnx" ]; then
	fail "${ID}" "Error exporting model"
	rm -rf "${UNZIPPED}"
	rm "${DEST}"
	exit 3
fi

MODELPATH=$(realpath "../models/${ID}/weights/best.onnx")

curl -s -X POST \
	-F "model@${MODELPATH}" \
	"${YOLO_WEB_URL}/job/${ID}"