add download big engine file to cdn manager

revise onnx export process
fixes
This commit is contained in:
Alex Bezdieniezhnykh
2025-04-24 16:23:39 +03:00
parent c4732cec8b
commit 26eb3b991d
9 changed files with 187 additions and 86 deletions
+27 -11
View File
@@ -1,13 +1,15 @@
import os
import shutil
from os import path, scandir, makedirs
from pathlib import Path
import random
import netron
import yaml
from ultralytics import YOLO
import constants
from azaion_api import Api, ApiCredentials
from api_client import ApiClient, ApiCredentials
from cdn_manager import CDNManager, CDNCredentials
from constants import datasets_dir, processed_images_dir
from security import Security
@@ -24,15 +26,21 @@ def export_rknn(model_path):
pass
def export_onnx(model_path):
def export_onnx(model_path, batch_size=4):
model = YOLO(model_path)
onnx_path = Path(model_path).stem + '.onnx'
if path.exists(onnx_path):
os.remove(onnx_path)
model.export(
format="onnx",
imgsz=1280,
batch=2,
batch=batch_size,
simplify=True,
nms=True)
return Path(model_path).stem + '.onnx'
nms=True,
device=0
)
return onnx_path
def export_tensorrt(model_path):
@@ -79,17 +87,25 @@ def upload_model(model_path: str, filename: str, size_small_in_kb: int=3):
key = Security.get_model_encryption_key()
model_encrypted = Security.encrypt_to(model_bytes, key)
part1_size = min(size_small_in_kb * 1024, int(0.9 * len(model_encrypted)))
part1_size = min(size_small_in_kb * 1024, int(0.3 * len(model_encrypted)))
model_part_small = model_encrypted[:part1_size] # slice bytes for part1
model_part_big = model_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, f'{filename}.big', model_part_big)
api = ApiClient(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder))
api = Api(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder))
api.upload_file(f'{filename}.small', model_part_small)
yaml_bytes = api.load_bytes(constants.CDN_CONFIG, '')
data = yaml.safe_load(yaml_bytes)
creds = CDNCredentials(data["host"],
data["downloader_access_key"],
data["downloader_access_secret"],
data["uploader_access_key"],
data["uploader_access_secret"])
cdn_manager = CDNManager(creds)
api.upload_file(f'{filename}.small', model_part_small, constants.MODELS_FOLDER)
cdn_manager.upload(constants.MODELS_FOLDER, f'{filename}.big', model_part_big)