Add core functionality for API client, CDN management, and data augmentation

- Introduced `ApiClient` for handling API interactions, including file uploads and downloads.
- Implemented `CDNManager` for managing CDN operations with AWS S3.
- Added `Augmentator` class for image augmentation, including bounding box corrections and transformations.
- Created utility functions for annotation conversion and dataset visualization.
- Established a new rules file for sound notifications during human input requests.

These additions enhance the system's capabilities for data handling and user interaction, laying the groundwork for future features.

Simplify autopilot state file to minimal current-step pointer; add execution safety rule to cursor-meta; remove Completed Steps/Key Decisions/Retry Log/Blockers from state template and all references.
This commit is contained in:
Oleksandr Bezdieniezhnykh
2026-03-28 00:12:54 +02:00
parent 142c6c4de8
commit c20018745b
31 changed files with 0 additions and 0 deletions
+148
View File
@@ -0,0 +1,148 @@
import re
import struct
import subprocess
from pathlib import Path
from typing import List, Tuple
import json
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
from inference.onnx_engine import InferenceEngine
# required for automatically initialize CUDA, do not remove.
import pycuda.autoinit
import pynvml
class TensorRTEngine(InferenceEngine):
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
def __init__(self, model_bytes: bytes, **kwargs):
try:
# metadata_len = struct.unpack("<I", model_bytes[:4])[0]
# try:
# self.metadata = json.loads(model_bytes[4:4 + metadata_len])
# self.class_names = self.metadata['names']
# print(f"Model metadata: {json.dumps(self.metadata, indent=2)}")
# except json.JSONDecodeError as err:
# print(f"Failed to parse metadata")
# return
# engine_data = model_bytes[4 + metadata_len:]
runtime = trt.Runtime(self.TRT_LOGGER)
self.engine = runtime.deserialize_cuda_engine(model_bytes)
if self.engine is None:
raise RuntimeError(f"Failed to load TensorRT engine!")
self.context = self.engine.create_execution_context()
# input
self.input_name = self.engine.get_tensor_name(0)
engine_input_shape = self.engine.get_tensor_shape(self.input_name)
if engine_input_shape[0] != -1:
self.batch_size = engine_input_shape[0]
self.input_shape = [
self.batch_size,
engine_input_shape[1], # Channels (usually fixed at 3 for RGB)
1280 if engine_input_shape[2] == -1 else engine_input_shape[2], # Height
1280 if engine_input_shape[3] == -1 else engine_input_shape[3] # Width
]
self.context.set_input_shape(self.input_name, self.input_shape)
input_size = trt.volume(self.input_shape) * np.dtype(np.float32).itemsize
self.d_input = cuda.mem_alloc(input_size)
# output
self.output_name = self.engine.get_tensor_name(1)
engine_output_shape = tuple(self.engine.get_tensor_shape(self.output_name))
self.output_shape = [
4 if self.input_shape[0] == -1 else self.input_shape[0], # by default, batch size is 4
300 if engine_output_shape[1] == -1 else engine_output_shape[1], # max detections number
6 if engine_output_shape[2] == -1 else engine_output_shape[2] # x1 y1 x2 y2 conf cls
]
self.h_output = cuda.pagelocked_empty(tuple(self.output_shape), dtype=np.float32)
self.d_output = cuda.mem_alloc(self.h_output.nbytes)
self.stream = cuda.Stream()
except Exception as e:
raise RuntimeError(f"Failed to initialize TensorRT engine: {str(e)}")
def get_input_shape(self) -> Tuple[int, int]:
return self.input_shape[2], self.input_shape[3]
def get_batch_size(self) -> int:
return self.batch_size
@staticmethod
def get_gpu_memory_bytes(device_id=0) -> int:
total_memory = None
try:
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
total_memory = mem_info.total
except pynvml.NVMLError:
total_memory = None
finally:
try:
pynvml.nvmlShutdown()
except pynvml.NVMLError:
pass
return 2 * 1024 * 1024 * 1024 if total_memory is None else total_memory # default 2 Gb
@staticmethod
def get_engine_filename(device_id=0) -> str | None:
try:
device = cuda.Device(device_id)
sm_count = device.multiprocessor_count
cc_major, cc_minor = device.compute_capability()
return f"azaion.cc_{cc_major}.{cc_minor}_sm_{sm_count}.engine"
except Exception:
return None
@staticmethod
def convert_from_onnx(onnx_model: bytes) -> bytes | None:
workspace_bytes = int(TensorRTEngine.get_gpu_memory_bytes() * 0.9)
explicit_batch_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(TensorRTEngine.TRT_LOGGER) as builder, \
builder.create_network(explicit_batch_flag) as network, \
trt.OnnxParser(network, TensorRTEngine.TRT_LOGGER) as parser, \
builder.create_builder_config() as config:
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes)
if not parser.parse(onnx_model):
return None
if builder.platform_has_fast_fp16:
print('Converting to supported fp16')
config.set_flag(trt.BuilderFlag.FP16)
else:
print('Converting to supported fp32. (fp16 is not supported)')
plan = builder.build_serialized_network(network, config)
if plan is None:
print('Conversion failed.')
return None
return bytes(plan)
def run(self, input_data: np.ndarray) -> List[np.ndarray]:
try:
cuda.memcpy_htod_async(self.d_input, input_data, self.stream)
self.context.set_tensor_address(self.input_name, int(self.d_input)) # input buffer
self.context.set_tensor_address(self.output_name, int(self.d_output)) # output buffer
self.context.execute_async_v3(stream_handle=self.stream.handle)
self.stream.synchronize()
# Fix: Remove the stream parameter from memcpy_dtoh
cuda.memcpy_dtoh(self.h_output, self.d_output)
output = self.h_output.reshape(self.output_shape)
return [output]
except Exception as e:
raise RuntimeError(f"Failed to run TensorRT inference: {str(e)}")