mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-23 05:36:35 +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.
54 lines
2.6 KiB
Markdown
54 lines
2.6 KiB
Markdown
# 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
|
|
None.
|