mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 21:26:36 +00:00
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.
This commit is contained in:
@@ -0,0 +1,53 @@
|
||||
# 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.
|
||||
Reference in New Issue
Block a user