train is ready

manual_run reuses train's export
add current_model to constants

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')

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):

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!')

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!')