109 lines
2.1 KiB
Bash
109 lines
2.1 KiB
Bash
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" --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 ..
|
|
|
|
if [ -f ./scripts/completed_mail.sh ]; then
|
|
bash ./scripts/completed_mail.sh "${JSON}"
|
|
else
|
|
echo "No email script, not sending"
|
|
fi
|