From 538dc8efa9719fe0291b9a6afc5eed788c7cad06 Mon Sep 17 00:00:00 2001
From: Alex Bezdieniezhnykh <zxsanny@gmail.com>
Date: Sat, 31 May 2025 18:41:10 +0300
Subject: [PATCH] train is ready manual_run reuses train's export add
 current_model to constants

---
 constants.py  |  9 +++------
 exports.py    |  1 -
 manual_run.py | 14 ++++----------
 train.py      | 27 +++++++++++----------------
 4 files changed, 18 insertions(+), 33 deletions(-)

diff --git a/constants.py b/constants.py
index 26b2c88..9a84aad 100644
--- a/constants.py
+++ b/constants.py
@@ -34,12 +34,9 @@ TXT_EXT = '.txt'
 
 OFFSET_FILE = 'offset.yaml'
 
-AI_ONNX_MODEL_FILE_BIG = "azaion.onnx.big"
-AI_ONNX_MODEL_FILE_SMALL = "azaion.onnx.small"
-
-AI_TENSOR_MODEL_FILE_BIG = "azaion.engine.big"
-AI_TENSOR_MODEL_FILE_SMALL = "azaion.engine.small"
-
 SMALL_SIZE_KB = 3
 CDN_CONFIG = 'cdn.yaml'
 MODELS_FOLDER = 'models'
+
+CURRENT_PT_MODEL = path.join(models_dir, f'{prefix[:-1]}.pt')
+CURRENT_ONNX_MODEL = path.join(models_dir, f'{prefix[:-1]}.onnx')
diff --git a/exports.py b/exports.py
index 57bb9b8..0794aff 100644
--- a/exports.py
+++ b/exports.py
@@ -40,7 +40,6 @@ def export_onnx(model_path, batch_size=4):
         nms=True,
         device=0
     )
-    return onnx_path
 
 
 def export_tensorrt(model_path):
diff --git a/manual_run.py b/manual_run.py
index a0c34fc..28ec2f2 100644
--- a/manual_run.py
+++ b/manual_run.py
@@ -2,7 +2,8 @@ import shutil
 from datetime import datetime
 from os import path
 
-from constants import models_dir, prefix, date_format, MODELS_FOLDER
+import train
+from constants import models_dir, prefix, date_format, MODELS_FOLDER, CURRENT_ONNX_MODEL
 from api_client import ApiClient
 from augmentation import Augmentator
 from exports import export_onnx
@@ -20,12 +21,5 @@ shutil.copytree(result_dir, model_dir, dirs_exist_ok=True)
 model_path = path.join(models_dir, f'{prefix[:-1]}.pt')
 shutil.copy(path.join(model_dir, 'weights', 'best.pt'), model_path)
 
-api_client = ApiClient()
-onnx_path = export_onnx(model_path)
-print(f'Conversion done: onnx path: {onnx_path}')
-
-with open(onnx_path, 'rb') as binary_file:
-    onnx_bytes = binary_file.read()
-
-key = Security.get_model_encryption_key()
-api_client.upload_big_small_resource(onnx_bytes, onnx_path, MODELS_FOLDER, key)
+train.export_current_model()
+print('success!')
\ No newline at end of file
diff --git a/train.py b/train.py
index 5a1c638..1850973 100644
--- a/train.py
+++ b/train.py
@@ -158,25 +158,20 @@ def train_dataset():
 
     model_dir = path.join(models_dir, today_folder)
     shutil.copytree(results.save_dir, model_dir)
-
-    model_path = path.join(models_dir, f'{prefix[:-1]}.pt')
-    shutil.copy(path.join(model_dir, 'weights', 'best.pt'), model_path)
-    return model_path
+    shutil.copy(path.join(model_dir, 'weights', 'best.pt'), constants.CURRENT_PT_MODEL)
 
 
-def validate(model_path):
-    model = YOLO(model_path)
-    print(model.val())
-
-
-if __name__ == '__main__':
-    model_path = train_dataset()
-    validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
-    onnx_path = export_onnx(model_path)
-
+def export_current_model():
+    export_onnx(constants.CURRENT_PT_MODEL)
     api_client = ApiClient()
-    with open(onnx_path, 'rb') as binary_file:
+    with open(constants.CURRENT_ONNX_MODEL, 'rb') as binary_file:
         onnx_bytes = binary_file.read()
 
     key = Security.get_model_encryption_key()
-    api_client.upload_big_small_resource(onnx_bytes, onnx_path, constants.MODELS_FOLDER, key)
+    api_client.upload_big_small_resource(onnx_bytes, 'azaion.onnx', constants.MODELS_FOLDER, key)
+
+
+if __name__ == '__main__':
+    train_dataset()
+    export_current_model()
+    print('success!')