upload model to cdn and api

switch to yolov11

train.py

@@ -1,3 +1,4 @@
+import io
 import os
 import random
 import shutil
@@ -6,7 +7,15 @@ from datetime import datetime
 from os import path, replace, listdir, makedirs, scandir
 from os.path import abspath
 from pathlib import Path
+from utils import Dotdict
+
+import yaml
 from ultralytics import YOLO
+
+import constants
+from azaion_api import ApiCredentials, Api
+from cdn_manager import CDNCredentials, CDNManager
+from security import Security
 from constants import (processed_images_dir,
                        processed_labels_dir,
                        annotation_classes,
@@ -145,6 +154,7 @@ def get_latest_model():
     last_model = sorted_dates[-1]
     return last_model['date'], last_model['path']
 
+
 def train_dataset(existing_date=None, from_scratch=False):
     latest_date, latest_model = get_latest_model()
 
@@ -156,7 +166,7 @@ def train_dataset(existing_date=None, from_scratch=False):
         cur_folder = today_folder
         cur_dataset = today_dataset
 
-    model_name = latest_model if latest_model is not None and path.isfile(latest_model) and not from_scratch else 'yolov8m.yaml'
+    model_name = latest_model if latest_model is not None and path.isfile(latest_model) and not from_scratch else 'yolo11m.yaml'
     print(f'Initial model: {model_name}')
     model = YOLO(model_name)
 
@@ -171,8 +181,10 @@ def train_dataset(existing_date=None, from_scratch=False):
     model_dir = path.join(models_dir, cur_folder)
     shutil.copytree(results.save_dir, model_dir)
 
-    shutil.copy(path.join(model_dir, 'weights', 'best.pt'), path.join(models_dir, f'{prefix[:-1]}.pt'))
+    model_path = path.join(models_dir, f'{prefix[:-1]}.pt')
+    shutil.copy(path.join(model_dir, 'weights', 'best.pt'), model_path)
     shutil.rmtree('runs')
+    return model_path
 
 
 def convert2rknn():
@@ -209,8 +221,38 @@ def validate(model_path):
     metrics = model.val()
     pass
 
+
+def upload_model(model_path: str):
+
+    # model = YOLO(model_path)
+    # model.export(format="onnx", imgsz=1280, nms=True, batch=4)
+    onnx_model = path.dirname(model_path) + Path(model_path).stem + '.onnx'
+
+    with open(onnx_model, 'rb') as f_in:
+        onnx_bytes = f_in.read()
+
+    key = Security.get_model_encryption_key()
+    onnx_encrypted = Security.encrypt_to(onnx_bytes, key)
+
+    part1_size = min(10 * 1024, int(0.9 * len(onnx_encrypted)))
+    onnx_part_small = onnx_encrypted[:part1_size]  # slice bytes for part1
+    onnx_part_big = onnx_encrypted[part1_size:]
+
+    with open(constants.CONFIG_FILE, "r") as f:
+        config_dict = yaml.safe_load(f)
+        d_config = Dotdict(config_dict)
+        cdn_c = Dotdict(d_config.cdn)
+        api_c = Dotdict(d_config.api)
+    cdn_manager = CDNManager(CDNCredentials(cdn_c.host, cdn_c.access_key, cdn_c.secret_key))
+    cdn_manager.upload(cdn_c.bucket, 'azaion.onnx.big', onnx_part_big)
+
+    api = Api(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder))
+    api.upload_file('azaion.onnx.small', onnx_part_small)
+
 if __name__ == '__main__':
-    train_dataset('2024-10-26', from_scratch=True)
-    validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
-    form_data_sample(500)
-    convert2rknn()
+    # model_path = train_dataset('2024-10-26', from_scratch=True)
+    # validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
+    # form_data_sample(500)
+    # convert2rknn()
+    model_path = 'azaion.pt'
+    upload_model(model_path)