mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 07:06:36 +00:00
Oleksandr Bezdieniezhnykh
142c6c4de8
- Replaced module-level path variables in constants.py with a structured Pydantic Config class. - Updated all relevant modules (train.py, augmentation.py, exports.py, dataset-visualiser.py, manual_run.py) to access paths through the new config structure. - Fixed bugs related to image processing and model saving. - Enhanced test infrastructure to accommodate the new configuration approach. This refactor improves code maintainability and clarity by centralizing configuration management.
119 lines
3.5 KiB
Python
119 lines
3.5 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 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=None):
|
|
if batch_size is None:
|
|
batch_size = constants.config.export.onnx_batch
|
|
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=constants.config.export.onnx_imgsz,
|
|
batch=batch_size,
|
|
simplify=True,
|
|
nms=True,
|
|
)
|
|
|
|
|
|
def export_coreml(model_path):
|
|
model = YOLO(model_path)
|
|
model.export(
|
|
format="coreml",
|
|
imgsz=constants.config.export.onnx_imgsz,
|
|
)
|
|
|
|
|
|
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(constants.config.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)
|