mirror of https://github.com/azaion/ai-training.git synced 2026-04-22 21:56:36 +00:00

Files

T

Oleksandr Bezdieniezhnykh 142c6c4de8 Refactor constants management to use Pydantic BaseModel for configuration

- 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.

2026-03-27 18:18:30 +02:00

2.6 KiB

Raw Blame History

Module: exports

Purpose

Model export utilities: converts trained YOLO .pt models to ONNX, TensorRT, and RKNN formats. Also handles encrypted model upload (split big/small pattern) and data sampling.

Public Interface

Function	Signature	Returns	Description
`export_rknn`	`(model_path: str)`	—	Exports YOLO model to RKNN format (RK3588 target), cleans up temp folder
`export_onnx`	`(model_path: str, batch_size: int = 4)`	—	Exports YOLO model to ONNX (1280px, NMS enabled, GPU device 0)
`export_tensorrt`	`(model_path: str)`	—	Exports YOLO model to TensorRT engine (batch=4, half precision, NMS)
`form_data_sample`	`(destination_path: str, size: int = 500, write_txt_log: bool = False)`	—	Creates a random sample of processed images
`show_model`	`(model: str = None)`	—	Opens model visualization in netron
`upload_model`	`(model_path: str, filename: str, size_small_in_kb: int = 3)`	—	Encrypts model, splits big/small, uploads to API + CDN

Internal Logic

export_onnx: Removes existing ONNX file if present, exports at 1280px with NMS baked in and simplification.
export_tensorrt: Uses YOLO's built-in TensorRT export (batch=4, FP16, NMS, simplify).
export_rknn: Exports to RKNN format targeting RK3588 SoC, moves result file and cleans temp directory.
upload_model: Encrypts with Security.get_model_encryption_key(), splits encrypted bytes at 30%/70% boundary (or size_small_in_kb * 1024), uploads small part to API, big part to CDN.
form_data_sample: Randomly shuffles processed images, copies first N to destination folder.

Dependencies

constants — directory paths, model paths, config file names
api_client — ApiClient, ApiCredentials for upload
cdn_manager — CDNManager, CDNCredentials for CDN upload
security — model encryption key, encrypt_to
utils — Dotdict for config access
ultralytics (external) — YOLO model
netron (external) — model visualization
yaml, os, shutil, random, pathlib (stdlib)

Consumers

train (export_tensorrt, upload_model, export_onnx)

Data Models

None.

Configuration

Reads config.yaml for API credentials (in upload_model), cdn.yaml for CDN credentials.

External Integrations

Ultralytics YOLO export pipeline
Netron model viewer
Azaion API + CDN for model upload

Security

Models are encrypted with AES-256-CBC before upload
Split storage (big on CDN, small on API) prevents single-point compromise

Tests