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 echo "Using medium weights [${MODEL}]" WEIGHTS="yolov5m.pt" elif [[ "${MODEL}" == "yolov5l_onnx" ]]; then echo "Using large weights [${MODEL}]" 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" --opset 12 --simplify --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") PTPATH=$(realpath "../models/${ID}/weights/best.pt") ZIPPATH="${ID}.zip" mkdir zip cp "${MODELPATH}" "zip/${NAME}.onnx" cp "${PTPATH}" "zip/${NAME}.pt" cd zip zip "../${ZIPPATH}" * cd .. rm -r zip/ curl -s -X POST \ -F "model=@${ZIPPATH}" \ "${YOLO_WEB_URL}/job/${ID}" rm "${ZIPPATH}" deactivate cd ..