From 226c972020b03fcb0de92dfa5c6fc4583c822eb4 Mon Sep 17 00:00:00 2001
From: mmcwilliams <hi@mmcwilliams.com>
Date: Wed, 29 Nov 2023 12:43:54 -0500
Subject: [PATCH] Use ONNX opset 12 (defaulted to 17) and simplify the model

---
 scripts/train.sh | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/scripts/train.sh b/scripts/train.sh
index d20b929..cf0b068 100644
--- a/scripts/train.sh
+++ b/scripts/train.sh
@@ -71,7 +71,7 @@ if [ ! -f "../models/${ID}/weights/best.pt" ]; then
 	exit 3
 fi
 
-python export.py --weights "../models/${ID}/weights/best.pt" --include onnx
+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"
@@ -81,11 +81,23 @@ if [ ! -f "../models/${ID}/weights/best.onnx" ]; then
 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=@${MODELPATH}" \
+	-F "model=@${ZIPPATH}" \
 	"${YOLO_WEB_URL}/job/${ID}"
 
+rm "${ZIPPATH}"
+
 deactivate
 cd ..