mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 08:46:36 +00:00
Alex Bezdieniezhnykh
112 lines
3.3 KiB
Python
112 lines
3.3 KiB
Python
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 api_client import ApiClient, ApiCredentials
|
|
from cdn_manager import CDNManager, CDNCredentials
|
|
from constants import datasets_dir, processed_images_dir
|
|
from security import Security
|
|
from utils import Dotdict
|
|
|
|
|
|
def export_rknn(model_path):
|
|
model = YOLO(model_path)
|
|
model.export(format="rknn", name="rk3588", simplify=True)
|
|
model_stem = Path(model_path).stem
|
|
folder_name = f'{model_stem}_rknn_model'
|
|
shutil.move(path.join(folder_name, f'{Path(model_path).stem}-rk3588.rknn'), f'{model_stem}.rknn')
|
|
shutil.rmtree(folder_name)
|
|
pass
|
|
|
|
|
|
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=batch_size,
|
|
simplify=True,
|
|
nms=True,
|
|
device=0
|
|
)
|
|
return onnx_path
|
|
|
|
|
|
def export_tensorrt(model_path):
|
|
YOLO(model_path).export(
|
|
format='engine',
|
|
batch=4,
|
|
half=True,
|
|
simplify=True,
|
|
nms=True
|
|
)
|
|
|
|
|
|
def form_data_sample(destination_path, size=500, write_txt_log=False):
|
|
images = []
|
|
with scandir(processed_images_dir) as imd:
|
|
for image_file in imd:
|
|
if not image_file.is_file():
|
|
continue
|
|
images.append(image_file)
|
|
print('shuffling images')
|
|
random.shuffle(images)
|
|
images = images[:size]
|
|
|
|
shutil.rmtree(destination_path, ignore_errors=True)
|
|
makedirs(destination_path, exist_ok=True)
|
|
|
|
lines = []
|
|
for image in images:
|
|
shutil.copy(image.path, path.join(destination_path, image.name))
|
|
lines.append(f'./{image.name}')
|
|
if write_txt_log:
|
|
with open(path.join(destination_path, 'azaion_subset.txt'), 'w', encoding='utf-8') as f:
|
|
f.writelines([f'{line}\n' for line in lines])
|
|
|
|
|
|
def show_model(model: str = None):
|
|
netron.start(model)
|
|
|
|
|
|
def upload_model(model_path: str, filename: str, size_small_in_kb: int=3):
|
|
with open(model_path, 'rb') as f_in:
|
|
model_bytes = f_in.read()
|
|
|
|
key = Security.get_model_encryption_key()
|
|
model_encrypted = Security.encrypt_to(model_bytes, key)
|
|
|
|
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)
|
|
api_c = Dotdict(d_config.api)
|
|
api = ApiClient(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder))
|
|
|
|
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)
|