From 522ad747178e7ff4b237ad64769a1b31e1d44aa1 Mon Sep 17 00:00:00 2001
From: mattmcw <matt@sixteenmillimeter.com>
Date: Fri, 30 Jun 2023 10:45:07 -0400
Subject: [PATCH] Caught typo in upload logic. Add export step to train

---
 scripts/claim.sh          |  6 +++++-
 scripts/install_yolov5.sh |  1 +
 scripts/train.sh          | 23 +++++++++++++++++++++++
 src/index.ts              |  4 ++--
 4 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/scripts/claim.sh b/scripts/claim.sh
index 2a9bf1f..31b862e 100644
--- a/scripts/claim.sh
+++ b/scripts/claim.sh
@@ -21,4 +21,8 @@ if [[ "${JSON}" != *"{"* ]]; then
 	exit 1
 fi
 
-bash ./scripts/train.sh "${JSON}"
\ No newline at end of file
+if [ -f ./scripts/claim_mail.sh ]; then
+	bash ./scripts/claim_mail.sh "${JSON}"
+fi
+
+bash ./scripts/train.sh "${JSON}"
diff --git a/scripts/install_yolov5.sh b/scripts/install_yolov5.sh
index 4ec96c6..c0d50a4 100644
--- a/scripts/install_yolov5.sh
+++ b/scripts/install_yolov5.sh
@@ -8,3 +8,4 @@ cd yolov5
 python -m venv env
 source env/bin/activate
 python -m pip install -r requirements.txt
+python -m pip install onnx
diff --git a/scripts/train.sh b/scripts/train.sh
index e98840d..ec50fa1 100644
--- a/scripts/train.sh
+++ b/scripts/train.sh
@@ -19,6 +19,8 @@ function fail () {
 		-d "{\"meta\":\"${META}\"}" "${URL}"
 }
 
+trap 'fail $ID' ERR
+
 JSON="${1}"
 
 ID=$(echo $JSON | jq -r '.id')
@@ -62,3 +64,24 @@ 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}"
diff --git a/src/index.ts b/src/index.ts
index 592ef3e..2de0e69 100755
--- a/src/index.ts
+++ b/src/index.ts
@@ -219,7 +219,7 @@ app.post('/', uploadZip.single('dataset'), async (req : Request, res : Response,
 
 	req.setTimeout(0);
 
-	if (typeof req.file === 'undefined' && req.file === null) {
+	if (typeof req.file === 'undefined' || req.file === null) {
 		console.error('No file in upload');
 		return next('ERROR: Please upload dataset as zip file');
 	}
@@ -273,7 +273,7 @@ app.post('/job/:id', uploadOnnx.single('model'), async (req : Request, res : Res
 
 	req.setTimeout(0);
 
-	if (typeof req.file === 'undefined' && req.file === null) {
+	if (typeof req.file === 'undefined' || req.file === null) {
 		console.error('No file in upload');
 		return next('ERROR: Please model as ONNX file');
 	}