mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 10:56: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:
Oleksandr Bezdieniezhnykh
@@ -0,0 +1,53 @@
|
||||
# Component: Core Infrastructure
|
||||
|
||||
## Overview
|
||||
Shared constants and utility classes that form the foundation for all other components. Provides path definitions, config file references, and helper data structures.
|
||||
|
||||
**Pattern**: Configuration constants + utility library
|
||||
**Upstream**: None (leaf component)
|
||||
**Downstream**: All other components
|
||||
|
||||
## Modules
|
||||
- `constants` — filesystem paths, config keys, thresholds
|
||||
- `utils` — Dotdict helper class
|
||||
|
||||
## Internal Interfaces
|
||||
|
||||
### constants (public symbols)
|
||||
All path/string constants — see module doc for full list. Key exports:
|
||||
- Directory paths: `data_dir`, `processed_dir`, `datasets_dir`, `models_dir` and their images/labels subdirectories
|
||||
- Config references: `CONFIG_FILE`, `CDN_CONFIG`, `OFFSET_FILE`
|
||||
- Model paths: `CURRENT_PT_MODEL`, `CURRENT_ONNX_MODEL`
|
||||
- Thresholds: `SMALL_SIZE_KB = 3`
|
||||
|
||||
### utils.Dotdict
|
||||
```python
|
||||
class Dotdict(dict):
|
||||
# Enables config.url instead of config["url"]
|
||||
```
|
||||
|
||||
## Data Access Patterns
|
||||
None — pure constants, no I/O.
|
||||
|
||||
## Implementation Details
|
||||
- All paths rooted at `/azaion/` — assumes a fixed deployment directory structure
|
||||
- No environment-variable override for any path — paths are entirely static
|
||||
|
||||
## Caveats
|
||||
- Hardcoded root `/azaion/` makes local development without that directory structure impossible
|
||||
- No `.env` or environment-based configuration override mechanism
|
||||
- `Dotdict.__getattr__` uses `dict.get` which returns `None` for missing keys instead of raising `AttributeError`
|
||||
|
||||
## Dependency Graph
|
||||
```mermaid
|
||||
graph TD
|
||||
constants --> api_client_comp[API & CDN]
|
||||
constants --> training_comp[Training]
|
||||
constants --> data_pipeline_comp[Data Pipeline]
|
||||
constants --> inference_comp[Inference]
|
||||
utils --> training_comp
|
||||
utils --> inference_comp
|
||||
```
|
||||
|
||||
## Logging Strategy
|
||||
None.
|
||||
@@ -0,0 +1,59 @@
|
||||
# Component: Security & Hardware Identity
|
||||
|
||||
## Overview
|
||||
Provides cryptographic operations (AES-256-CBC encryption/decryption) and hardware fingerprinting. Used for protecting model files in transit and at rest, and for binding API encryption keys to specific machines.
|
||||
|
||||
**Pattern**: Utility/service library (static methods)
|
||||
**Upstream**: None (leaf component)
|
||||
**Downstream**: API & CDN, Training, Inference
|
||||
|
||||
## Modules
|
||||
- `security` — AES encryption, key derivation (SHA-384), hardcoded model key
|
||||
- `hardware_service` — cross-platform hardware info collection (CPU, GPU, RAM, drive serial)
|
||||
|
||||
## Internal Interfaces
|
||||
|
||||
### Security (static methods)
|
||||
```python
|
||||
Security.encrypt_to(input_bytes: bytes, key: str) -> bytes
|
||||
Security.decrypt_to(ciphertext_with_iv: bytes, key: str) -> bytes
|
||||
Security.calc_hash(key: str) -> str
|
||||
Security.get_hw_hash(hardware: str) -> str
|
||||
Security.get_api_encryption_key(creds, hardware_hash: str) -> str
|
||||
Security.get_model_encryption_key() -> str
|
||||
```
|
||||
|
||||
### hardware_service
|
||||
```python
|
||||
get_hardware_info() -> str
|
||||
```
|
||||
|
||||
## Data Access Patterns
|
||||
- `hardware_service` executes shell commands to query OS/hardware info
|
||||
- `security` performs in-memory cryptographic operations only
|
||||
|
||||
## Implementation Details
|
||||
- **Encryption**: AES-256-CBC. Key = SHA-256(key_string). IV = 16 random bytes prepended to ciphertext. PKCS7 padding.
|
||||
- **Key derivation hierarchy**:
|
||||
1. `get_model_encryption_key()` → hardcoded secret → SHA-384 → base64
|
||||
2. `get_hw_hash(hardware_string)` → salted hardware string → SHA-384 → base64
|
||||
3. `get_api_encryption_key(creds, hw_hash)` → email+password+hw_hash+salt → SHA-384 → base64
|
||||
- **Hardware fingerprint format**: `CPU: {cpu}. GPU: {gpu}. Memory: {memory}. DriveSerial: {serial}`
|
||||
|
||||
## Caveats
|
||||
- **Hardcoded model encryption key** in `get_model_encryption_key()` — anyone with source code access can derive the key
|
||||
- **Shell command injection risk**: `hardware_service` uses `shell=True` subprocess — safe since no user input is involved, but fragile
|
||||
- **PKCS7 unpadding** in `decrypt_to` uses manual check instead of the `cryptography` library's unpadder — potential padding oracle if error handling is observed
|
||||
- `BUFFER_SIZE` constant declared but unused in security.py
|
||||
|
||||
## Dependency Graph
|
||||
```mermaid
|
||||
graph TD
|
||||
hardware_service --> api_client[API & CDN: api_client]
|
||||
security --> api_client
|
||||
security --> training[Training]
|
||||
security --> inference[Inference: start_inference]
|
||||
```
|
||||
|
||||
## Logging Strategy
|
||||
None — operations are silent except for exceptions.
|
||||
@@ -0,0 +1,90 @@
|
||||
# Component: API & CDN Client
|
||||
|
||||
## Overview
|
||||
Communication layer for the Azaion backend API and S3-compatible CDN. Handles authentication, encrypted file transfer, and the split-resource pattern for secure model distribution.
|
||||
|
||||
**Pattern**: Client library with split-storage resource management
|
||||
**Upstream**: Core (constants), Security (encryption, hardware identity)
|
||||
**Downstream**: Training, Inference, Exports
|
||||
|
||||
## Modules
|
||||
- `api_client` — REST client for Azaion API, JWT auth, encrypted resource download/upload, split big/small pattern
|
||||
- `cdn_manager` — boto3 S3 client with separate read/write credentials
|
||||
|
||||
## Internal Interfaces
|
||||
|
||||
### CDNCredentials
|
||||
```python
|
||||
CDNCredentials(host, downloader_access_key, downloader_access_secret, uploader_access_key, uploader_access_secret)
|
||||
```
|
||||
|
||||
### CDNManager
|
||||
```python
|
||||
CDNManager(credentials: CDNCredentials)
|
||||
CDNManager.upload(bucket: str, filename: str, file_bytes: bytearray) -> bool
|
||||
CDNManager.download(bucket: str, filename: str) -> bool
|
||||
```
|
||||
|
||||
### ApiCredentials
|
||||
```python
|
||||
ApiCredentials(url, email, password)
|
||||
```
|
||||
|
||||
### ApiClient
|
||||
```python
|
||||
ApiClient()
|
||||
ApiClient.login() -> None
|
||||
ApiClient.upload_file(filename: str, file_bytes: bytearray, folder: str) -> None
|
||||
ApiClient.load_bytes(filename: str, folder: str) -> bytes
|
||||
ApiClient.load_big_small_resource(resource_name: str, folder: str, key: str) -> bytes
|
||||
ApiClient.upload_big_small_resource(resource: bytes, resource_name: str, folder: str, key: str) -> None
|
||||
```
|
||||
|
||||
## External API Specification
|
||||
|
||||
### Azaion REST API (consumed)
|
||||
| Endpoint | Method | Auth | Description |
|
||||
|----------|--------|------|-------------|
|
||||
| `/login` | POST | None (returns JWT) | `{"email": ..., "password": ...}` → `{"token": ...}` |
|
||||
| `/resources/{folder}` | POST | Bearer JWT | Multipart file upload |
|
||||
| `/resources/get/{folder}` | POST | Bearer JWT | Download encrypted resource (sends hardware info in body) |
|
||||
|
||||
### S3-compatible CDN
|
||||
| Operation | Description |
|
||||
|-----------|-------------|
|
||||
| `upload_fileobj` | Upload bytes to S3 bucket |
|
||||
| `download_file` | Download file from S3 bucket to disk |
|
||||
|
||||
## Data Access Patterns
|
||||
- API Client reads `config.yaml` on init for API credentials
|
||||
- CDN credentials loaded by API Client from encrypted `cdn.yaml` (downloaded from API)
|
||||
- Split resources: big part stored locally + CDN, small part on API server
|
||||
|
||||
## Implementation Details
|
||||
- **JWT auto-refresh**: On 401/403 response, automatically re-authenticates and retries
|
||||
- **Split-resource pattern**: Encrypts data → splits at ~20% (SMALL_SIZE_KB * 1024 min) boundary → small part to API, big part to CDN. Neither part alone can reconstruct the original.
|
||||
- **CDN credential isolation**: Separate S3 access keys for upload vs download (least-privilege)
|
||||
- **CDN self-bootstrap**: `cdn.yaml` credentials are themselves encrypted and downloaded from the API during ApiClient init
|
||||
|
||||
## Caveats
|
||||
- Credentials hardcoded in `config.yaml` and `cdn.yaml` — not using environment variables or secrets manager
|
||||
- `cdn_manager.download()` saves to current working directory with the same filename
|
||||
- No retry logic beyond JWT refresh (no exponential backoff, no connection retry)
|
||||
- `CDNManager` imports `sys`, `yaml`, `os` but doesn't use them
|
||||
|
||||
## Dependency Graph
|
||||
```mermaid
|
||||
graph TD
|
||||
constants --> api_client
|
||||
security --> api_client
|
||||
hardware_service --> api_client
|
||||
cdn_manager --> api_client
|
||||
api_client --> exports
|
||||
api_client --> train
|
||||
api_client --> start_inference
|
||||
cdn_manager --> exports
|
||||
cdn_manager --> train
|
||||
```
|
||||
|
||||
## Logging Strategy
|
||||
Print statements for upload/download confirmations and errors. No structured logging.
|
||||
@@ -0,0 +1,61 @@
|
||||
# Component: Data Models
|
||||
|
||||
## Overview
|
||||
Shared data transfer objects for the training pipeline: annotation class definitions (with weather modes) and image+label containers for visualization and augmentation.
|
||||
|
||||
**Pattern**: Plain data classes / value objects
|
||||
**Upstream**: None (leaf)
|
||||
**Downstream**: Data Pipeline (augmentation, dataset-visualiser), Training (YAML generation)
|
||||
|
||||
## Modules
|
||||
- `dto/annotationClass` — AnnotationClass, WeatherMode enum, classes.json reader
|
||||
- `dto/imageLabel` — ImageLabel container with bbox visualization
|
||||
|
||||
## Internal Interfaces
|
||||
|
||||
### WeatherMode (Enum)
|
||||
| Member | Value | Description |
|
||||
|--------|-------|-------------|
|
||||
| Norm | 0 | Normal weather |
|
||||
| Wint | 20 | Winter |
|
||||
| Night | 40 | Night |
|
||||
|
||||
### AnnotationClass
|
||||
```python
|
||||
AnnotationClass(id: int, name: str, color: str)
|
||||
AnnotationClass.read_json() -> dict[int, AnnotationClass] # static
|
||||
AnnotationClass.color_tuple -> tuple # property, RGB ints
|
||||
```
|
||||
|
||||
### ImageLabel
|
||||
```python
|
||||
ImageLabel(image_path: str, image: np.ndarray, labels_path: str, labels: list)
|
||||
ImageLabel.visualize(annotation_classes: dict) -> None
|
||||
```
|
||||
|
||||
## Data Access Patterns
|
||||
- `AnnotationClass.read_json()` reads `classes.json` from project root (relative to `dto/` parent)
|
||||
- `ImageLabel.visualize()` renders to matplotlib window (no disk I/O)
|
||||
|
||||
## Implementation Details
|
||||
- 17 base annotation classes × 3 weather modes = 51 classes with offset IDs (0–16, 20–36, 40–56)
|
||||
- System reserves 80 class slots (DEFAULT_CLASS_NUM in train.py)
|
||||
- YOLO label format: [x_center, y_center, width, height, class_id] — all normalized 0–1
|
||||
- `color_tuple` parsing strips first 3 chars (assumes "#ff" prefix format) — fragile if color format changes
|
||||
|
||||
## Caveats
|
||||
- `AnnotationClass` duplicated in 3 locations (dto, inference/dto, annotation-queue/annotation_queue_dto) with slight differences
|
||||
- `color_tuple` property has a non-obvious parsing approach that may break on different color string formats
|
||||
- Empty files: `dto/annotation_bulk_message.py` and `dto/annotation_message.py` suggest planned but unimplemented DTOs
|
||||
|
||||
## Dependency Graph
|
||||
```mermaid
|
||||
graph TD
|
||||
dto_annotationClass[dto/annotationClass] --> train
|
||||
dto_annotationClass --> dataset-visualiser
|
||||
dto_imageLabel[dto/imageLabel] --> augmentation
|
||||
dto_imageLabel --> dataset-visualiser
|
||||
```
|
||||
|
||||
## Logging Strategy
|
||||
None.
|
||||
@@ -0,0 +1,74 @@
|
||||
# Component: Data Pipeline
|
||||
|
||||
## Overview
|
||||
Tools for preparing and managing annotation data: augmentation of training images, format conversion from external annotation systems, and visual inspection of annotated datasets.
|
||||
|
||||
**Pattern**: Batch processing tools (standalone scripts + library)
|
||||
**Upstream**: Core (constants), Data Models (ImageLabel, AnnotationClass)
|
||||
**Downstream**: Training (augmented images feed into dataset formation)
|
||||
|
||||
## Modules
|
||||
- `augmentation` — image augmentation pipeline (albumentations)
|
||||
- `convert-annotations` — Pascal VOC / oriented bbox → YOLO format converter
|
||||
- `dataset-visualiser` — interactive annotation visualization tool
|
||||
|
||||
## Internal Interfaces
|
||||
|
||||
### Augmentator
|
||||
```python
|
||||
Augmentator()
|
||||
Augmentator.augment_annotations(from_scratch: bool = False) -> None
|
||||
Augmentator.augment_inner(img_ann: ImageLabel) -> list[ImageLabel]
|
||||
Augmentator.correct_bboxes(labels) -> list
|
||||
Augmentator.read_labels(labels_path) -> list[list]
|
||||
```
|
||||
|
||||
### convert-annotations (functions)
|
||||
```python
|
||||
convert(folder, dest_folder, read_annotations, ann_format) -> None
|
||||
minmax2yolo(width, height, xmin, xmax, ymin, ymax) -> tuple
|
||||
read_pascal_voc(width, height, s: str) -> list[str]
|
||||
read_bbox_oriented(width, height, s: str) -> list[str]
|
||||
```
|
||||
|
||||
### dataset-visualiser (functions)
|
||||
```python
|
||||
visualise_dataset() -> None
|
||||
visualise_processed_folder() -> None
|
||||
```
|
||||
|
||||
## Data Access Patterns
|
||||
- **Augmentation**: Reads from `/azaion/data/images/` + `/azaion/data/labels/`, writes to `/azaion/data-processed/images/` + `/azaion/data-processed/labels/`
|
||||
- **Conversion**: Reads from user-specified source folder, writes to destination folder
|
||||
- **Visualiser**: Reads from datasets or processed folder, renders to matplotlib window
|
||||
|
||||
## Implementation Details
|
||||
- **Augmentation pipeline**: Per image → 1 original copy + 7 augmented variants (8× data expansion)
|
||||
- HorizontalFlip (60%), BrightnessContrast (40%), Affine (80%), MotionBlur (10%), HueSaturation (40%)
|
||||
- Bbox correction clips outside-boundary boxes, removes boxes < 1% of image
|
||||
- Incremental: skips already-processed images
|
||||
- Continuous mode: infinite loop with 5-minute sleep between rounds
|
||||
- Concurrent: ThreadPoolExecutor for parallel image processing
|
||||
- **Format conversion**: Pluggable reader pattern — `convert()` accepts any reader function that maps (width, height, text) → YOLO lines
|
||||
- **Visualiser**: Interactive (waits for keypress) — developer debugging tool
|
||||
|
||||
## Caveats
|
||||
- `dataset-visualiser` imports from `preprocessing` module which does not exist — broken import
|
||||
- `dataset-visualiser` has hardcoded dataset date (`2024-06-18`) and start index (35247)
|
||||
- `convert-annotations` hardcodes class mappings (Truck=1, Car/Taxi=2) — not configurable
|
||||
- Augmentation parameters are hardcoded, not configurable via config file
|
||||
- Augmentation `total_to_process` attribute referenced in `augment_annotation` but never set (uses `total_images_to_process`)
|
||||
|
||||
## Dependency Graph
|
||||
```mermaid
|
||||
graph TD
|
||||
constants --> augmentation
|
||||
dto_imageLabel[dto/imageLabel] --> augmentation
|
||||
constants --> dataset-visualiser
|
||||
dto_annotationClass[dto/annotationClass] --> dataset-visualiser
|
||||
dto_imageLabel --> dataset-visualiser
|
||||
augmentation --> manual_run
|
||||
```
|
||||
|
||||
## Logging Strategy
|
||||
Print statements for progress tracking (processed count, errors). No structured logging.
|
||||
@@ -0,0 +1,87 @@
|
||||
# Component: Training Pipeline
|
||||
|
||||
## Overview
|
||||
End-to-end YOLOv11 object detection training workflow: dataset formation from augmented annotations, model training, multi-format export (ONNX, TensorRT, RKNN), and encrypted model upload.
|
||||
|
||||
**Pattern**: Pipeline / orchestrator
|
||||
**Upstream**: Core, Security, API & CDN, Data Models, Data Pipeline (augmented images)
|
||||
**Downstream**: None (produces trained models consumed externally)
|
||||
|
||||
## Modules
|
||||
- `train` — main pipeline: dataset formation → YOLO training → export → upload
|
||||
- `exports` — model format conversion (ONNX, TensorRT, RKNN) + upload utilities
|
||||
- `manual_run` — ad-hoc developer script for selective pipeline steps
|
||||
|
||||
## Internal Interfaces
|
||||
|
||||
### train
|
||||
```python
|
||||
form_dataset() -> None
|
||||
copy_annotations(images, folder: str) -> None
|
||||
check_label(label_path: str) -> bool
|
||||
create_yaml() -> None
|
||||
resume_training(last_pt_path: str) -> None
|
||||
train_dataset() -> None
|
||||
export_current_model() -> None
|
||||
```
|
||||
|
||||
### exports
|
||||
```python
|
||||
export_rknn(model_path: str) -> None
|
||||
export_onnx(model_path: str, batch_size: int = 4) -> None
|
||||
export_tensorrt(model_path: str) -> None
|
||||
form_data_sample(destination_path: str, size: int = 500, write_txt_log: bool = False) -> None
|
||||
show_model(model: str = None) -> None
|
||||
upload_model(model_path: str, filename: str, size_small_in_kb: int = 3) -> None
|
||||
```
|
||||
|
||||
## Data Access Patterns
|
||||
- **Input**: Reads augmented images from `/azaion/data-processed/images/` + labels
|
||||
- **Dataset output**: Creates dated dataset at `/azaion/datasets/azaion-{YYYY-MM-DD}/` with train/valid/test splits
|
||||
- **Model output**: Saves trained models to `/azaion/models/azaion-{YYYY-MM-DD}/`, copies best.pt to `/azaion/models/azaion.pt`
|
||||
- **Upload**: Encrypted model uploaded as split big/small to CDN + API
|
||||
- **Corrupted data**: Invalid labels moved to `/azaion/data-corrupted/`
|
||||
|
||||
## Implementation Details
|
||||
- **Dataset split**: 70% train / 20% valid / 10% test (random shuffle)
|
||||
- **Label validation**: `check_label()` verifies all YOLO coordinates are ≤ 1.0
|
||||
- **YAML generation**: Writes `data.yaml` with 80 class names (17 actual from classes.json × 3 weather modes, rest as placeholders)
|
||||
- **Training config**: YOLOv11 medium (`yolo11m.yaml`), epochs=120, batch=11 (tuned for 24GB VRAM), imgsz=1280, save_period=1, workers=24
|
||||
- **Post-training**: Removes intermediate epoch checkpoints, keeps only `best.pt`
|
||||
- **Export chain**: `.pt` → ONNX (1280px, batch=4, NMS) → encrypted → split → upload
|
||||
- **TensorRT export**: batch=4, FP16, NMS, simplify
|
||||
- **RKNN export**: targets RK3588 SoC (OrangePi5)
|
||||
- **Concurrent file copying**: ThreadPoolExecutor for parallel image/label copying during dataset formation
|
||||
- **`__main__`** in `train.py`: `train_dataset()` → `export_current_model()`
|
||||
|
||||
## Caveats
|
||||
- Training hyperparameters are hardcoded (not configurable via config file)
|
||||
- `old_images_percentage = 75` declared but unused
|
||||
- `train.py` imports `subprocess`, `sleep` but doesn't use them
|
||||
- `train.py` imports `OnnxEngine` but doesn't use it
|
||||
- `exports.upload_model()` creates `ApiClient` with different constructor signature than the one in `api_client.py` — likely stale code
|
||||
- `copy_annotations` uses a global `total_files_copied` counter with a local `copied` variable that stays at 0 — reporting bug
|
||||
- `resume_training` references `yaml` (the module) instead of a YAML file path in the `data` parameter
|
||||
|
||||
## Dependency Graph
|
||||
```mermaid
|
||||
graph TD
|
||||
constants --> train
|
||||
constants --> exports
|
||||
api_client --> train
|
||||
api_client --> exports
|
||||
cdn_manager --> train
|
||||
cdn_manager --> exports
|
||||
security --> train
|
||||
security --> exports
|
||||
utils --> train
|
||||
utils --> exports
|
||||
dto_annotationClass[dto/annotationClass] --> train
|
||||
inference_onnx[inference/onnx_engine] --> train
|
||||
exports --> train
|
||||
train --> manual_run
|
||||
augmentation --> manual_run
|
||||
```
|
||||
|
||||
## Logging Strategy
|
||||
Print statements for progress (file count, shuffling status, training results). No structured logging.
|
||||
@@ -0,0 +1,85 @@
|
||||
# Component: Inference Engine
|
||||
|
||||
## Overview
|
||||
Real-time object detection inference subsystem supporting ONNX Runtime and TensorRT backends. Processes video streams with batched inference, custom NMS, and live visualization.
|
||||
|
||||
**Pattern**: Strategy pattern (InferenceEngine ABC) + pipeline orchestrator
|
||||
**Upstream**: Core, Security, API & CDN (for model download)
|
||||
**Downstream**: None (end-user facing — processes video input)
|
||||
|
||||
## Modules
|
||||
- `inference/dto` — Detection, Annotation, AnnotationClass data classes
|
||||
- `inference/onnx_engine` — InferenceEngine ABC + OnnxEngine implementation
|
||||
- `inference/tensorrt_engine` — TensorRTEngine implementation with CUDA memory management + ONNX converter
|
||||
- `inference/inference` — Video processing pipeline (preprocess → infer → postprocess → draw)
|
||||
- `start_inference` — Entry point: downloads model, initializes engine, runs on video
|
||||
|
||||
## Internal Interfaces
|
||||
|
||||
### InferenceEngine (ABC)
|
||||
```python
|
||||
InferenceEngine.__init__(model_path: str, batch_size: int = 1, **kwargs)
|
||||
InferenceEngine.get_input_shape() -> Tuple[int, int]
|
||||
InferenceEngine.get_batch_size() -> int
|
||||
InferenceEngine.run(input_data: np.ndarray) -> List[np.ndarray]
|
||||
```
|
||||
|
||||
### OnnxEngine (extends InferenceEngine)
|
||||
Constructor takes `model_bytes` (not path). Uses CUDAExecutionProvider + CPUExecutionProvider.
|
||||
|
||||
### TensorRTEngine (extends InferenceEngine)
|
||||
Constructor takes `model_bytes: bytes`. Additional static methods:
|
||||
```python
|
||||
TensorRTEngine.get_gpu_memory_bytes(device_id=0) -> int
|
||||
TensorRTEngine.get_engine_filename(device_id=0) -> str | None
|
||||
TensorRTEngine.convert_from_onnx(onnx_model: bytes) -> bytes | None
|
||||
```
|
||||
|
||||
### Inference
|
||||
```python
|
||||
Inference(engine: InferenceEngine, confidence_threshold, iou_threshold)
|
||||
Inference.preprocess(frames: list) -> np.ndarray
|
||||
Inference.postprocess(batch_frames, batch_timestamps, output) -> list[Annotation]
|
||||
Inference.process(video: str) -> None
|
||||
Inference.draw(annotation: Annotation) -> None
|
||||
Inference.remove_overlapping_detections(detections) -> list[Detection]
|
||||
```
|
||||
|
||||
## Data Access Patterns
|
||||
- Model bytes loaded by caller (start_inference via ApiClient.load_big_small_resource)
|
||||
- Video input via cv2.VideoCapture (file path)
|
||||
- No disk writes during inference
|
||||
|
||||
## Implementation Details
|
||||
- **Video processing**: Every 4th frame processed (25% frame sampling), batched to engine batch size
|
||||
- **Preprocessing**: cv2.dnn.blobFromImage (1/255 scale, model input size, BGR→RGB)
|
||||
- **Postprocessing**: Raw detections filtered by confidence, coordinates normalized to [0,1], custom NMS applied
|
||||
- **Custom NMS**: Pairwise IoU comparison. Keeps higher confidence; ties broken by lower class ID.
|
||||
- **TensorRT**: Async CUDA execution (memcpy_htod_async → execute_async_v3 → synchronize → memcpy_dtoh)
|
||||
- **TensorRT shapes**: Default 1280×1280 input, 300 max detections, 6 values per detection (x1,y1,x2,y2,conf,cls)
|
||||
- **ONNX conversion**: TensorRT builder with 90% GPU memory workspace, FP16 if supported
|
||||
- **Engine filename**: GPU-architecture-specific: `azaion.cc_{major}.{minor}_sm_{sm_count}.engine`
|
||||
- **start_inference flow**: ApiClient → load encrypted TensorRT model (big/small split) → decrypt → TensorRTEngine → Inference.process()
|
||||
|
||||
## Caveats
|
||||
- `start_inference.get_engine_filename()` duplicates `TensorRTEngine.get_engine_filename()`
|
||||
- Video path hardcoded in `start_inference` (`tests/ForAI_test.mp4`)
|
||||
- `inference/dto` has its own AnnotationClass — duplicated from `dto/annotationClass`
|
||||
- cv2.imshow display requires a GUI environment — won't work headless
|
||||
- TensorRT `batch_size` attribute used before assignment if engine input shape has dynamic batch — potential NameError
|
||||
|
||||
## Dependency Graph
|
||||
```mermaid
|
||||
graph TD
|
||||
inference_dto[inference/dto] --> inference_inference[inference/inference]
|
||||
inference_onnx[inference/onnx_engine] --> inference_inference
|
||||
inference_onnx --> inference_trt[inference/tensorrt_engine]
|
||||
inference_trt --> start_inference
|
||||
inference_inference --> start_inference
|
||||
constants --> start_inference
|
||||
api_client --> start_inference
|
||||
security --> start_inference
|
||||
```
|
||||
|
||||
## Logging Strategy
|
||||
Print statements for metadata, download progress, timing. cv2.imshow for visual output.
|
||||
@@ -0,0 +1,71 @@
|
||||
# Component: Annotation Queue Service
|
||||
|
||||
## Overview
|
||||
Self-contained async service that consumes annotation CRUD events from a RabbitMQ Streams queue and persists images + labels to the filesystem. Operates independently from the training pipeline.
|
||||
|
||||
**Pattern**: Message-driven event handler / consumer service
|
||||
**Upstream**: External RabbitMQ Streams queue (Azaion platform)
|
||||
**Downstream**: Data Pipeline (files written become input for augmentation)
|
||||
|
||||
## Modules
|
||||
- `annotation-queue/annotation_queue_dto` — message DTOs (AnnotationMessage, AnnotationBulkMessage, AnnotationStatus, Detection, etc.)
|
||||
- `annotation-queue/annotation_queue_handler` — async queue consumer with message routing and file management
|
||||
|
||||
## Internal Interfaces
|
||||
|
||||
### AnnotationQueueHandler
|
||||
```python
|
||||
AnnotationQueueHandler()
|
||||
AnnotationQueueHandler.start() -> async
|
||||
AnnotationQueueHandler.on_message(message: AMQPMessage, context: MessageContext) -> None
|
||||
AnnotationQueueHandler.save_annotation(ann: AnnotationMessage) -> None
|
||||
AnnotationQueueHandler.validate(msg: AnnotationBulkMessage) -> None
|
||||
AnnotationQueueHandler.delete(msg: AnnotationBulkMessage) -> None
|
||||
```
|
||||
|
||||
### Key DTOs
|
||||
```python
|
||||
AnnotationMessage(msgpack_bytes) # Full annotation with image + detections
|
||||
AnnotationBulkMessage(msgpack_bytes) # Bulk validate/delete
|
||||
AnnotationStatus: Created(10), Edited(20), Validated(30), Deleted(40)
|
||||
RoleEnum: Operator(10), Validator(20), CompanionPC(30), Admin(40), ApiAdmin(1000)
|
||||
```
|
||||
|
||||
## Data Access Patterns
|
||||
- **Queue**: Consumes from RabbitMQ Streams queue `azaion-annotations` using rstream library
|
||||
- **Offset persistence**: `offset.yaml` tracks last processed message offset for resume
|
||||
- **Filesystem writes**:
|
||||
- Validated annotations → `{root}/data/images/` + `{root}/data/labels/`
|
||||
- Unvalidated (seed) → `{root}/data-seed/images/` + `{root}/data-seed/labels/`
|
||||
- Deleted → `{root}/data_deleted/images/` + `{root}/data_deleted/labels/`
|
||||
|
||||
## Implementation Details
|
||||
- **Message routing**: Based on `AnnotationStatus` from AMQP application properties:
|
||||
- Created/Edited → save label + optionally image; validator role writes to data, operator to seed
|
||||
- Validated (bulk) → move from seed to data
|
||||
- Deleted (bulk) → move to deleted directory
|
||||
- **Role-based logic**: `RoleEnum.is_validator()` returns True for Validator, Admin, ApiAdmin — these roles write directly to validated data directory
|
||||
- **Serialization**: Messages are msgpack-encoded with positional integer keys. Detections are embedded as a JSON string within the msgpack payload.
|
||||
- **Offset tracking**: After each successfully processed message, offset is persisted to `offset.yaml` (survives restarts)
|
||||
- **Logging**: TimedRotatingFileHandler with daily rotation, 7-day retention, writes to `logs/` directory
|
||||
- **Separate dependencies**: Own `requirements.txt` (pyyaml, msgpack, rstream only)
|
||||
- **Own config.yaml**: Points to test directories by default (`data-test`, `data-test-seed`)
|
||||
|
||||
## Caveats
|
||||
- Credentials hardcoded in `config.yaml` (queue host, user, password)
|
||||
- AnnotationClass duplicated (third copy) with slight differences from dto/ version
|
||||
- No reconnection logic for queue disconnections
|
||||
- No dead-letter queue or message retry on processing failures
|
||||
- `save_annotation` writes empty label files when detections list has no newline separators between entries
|
||||
- The annotation-queue `config.yaml` uses different directory names (`data-test` vs `data`) than the main `config.yaml` — likely a test vs production configuration issue
|
||||
|
||||
## Dependency Graph
|
||||
```mermaid
|
||||
graph TD
|
||||
annotation_queue_dto --> annotation_queue_handler
|
||||
rstream_ext[rstream library] --> annotation_queue_handler
|
||||
msgpack_ext[msgpack library] --> annotation_queue_dto
|
||||
```
|
||||
|
||||
## Logging Strategy
|
||||
`logging` module with TimedRotatingFileHandler. Format: `HH:MM:SS|message`. Daily rotation, 7-day retention. Also outputs to stdout.
|
||||
Reference in New Issue
Block a user