diff --git a/azaion_api.py b/api_client.py similarity index 54% rename from azaion_api.py rename to api_client.py index 14b4f0c..8c1a01b 100644 --- a/azaion_api.py +++ b/api_client.py @@ -1,26 +1,43 @@ import io import json from http import HTTPStatus +from os import path import requests +import yaml import constants +from cdn_manager import CDNCredentials, CDNManager from hardware_service import get_hardware_info from security import Security class ApiCredentials: - def __init__(self, url, email, password, folder): + def __init__(self, url, email, password): self.url = url self.email = email self.password = password - self.folder = folder -class Api: - def __init__(self, credentials): +class ApiClient: + def __init__(self): self.token = None - self.credentials = credentials + with open(constants.CONFIG_FILE, "r") as f: + config_dict = yaml.safe_load(f) + api_c = config_dict["api"] + self.credentials = ApiCredentials(api_c["url"], + api_c["email"], + api_c["password"]) + + yaml_bytes = self.load_bytes(constants.CDN_CONFIG, '') + data = yaml.safe_load(io.BytesIO(yaml_bytes)) + creds = CDNCredentials(data["host"], + data["downloader_access_key"], + data["downloader_access_secret"], + data["uploader_access_key"], + data["uploader_access_secret"]) + + self.cdn_manager = CDNManager(creds) def login(self): response = requests.post(f'{self.credentials.url}/login', @@ -29,8 +46,7 @@ class Api: token = response.json()["token"] self.token = token - def upload_file(self, filename: str, file_bytes: bytearray): - folder = self.credentials.folder + def upload_file(self, filename: str, file_bytes: bytearray, folder): if self.token is None: self.login() url = f"{self.credentials.url}/resources/{folder}" @@ -43,9 +59,7 @@ class Api: except Exception as e: print(f"Upload fail: {e}") - def load_bytes(self, filename, folder = None): - folder = folder or self.credentials.folder - + def load_bytes(self, filename, folder): hardware = get_hardware_info() if self.token is None: @@ -82,13 +96,30 @@ class Api: print(f'Downloaded file: {filename}, {len(data)} bytes') return data - def load_resource(self, big_part, small_part): + def load_big_small_resource(self, resource_name, folder, key): + big_part = path.join(folder, f'{resource_name}.big') + small_part = f'{resource_name}.small' + with open(big_part, 'rb') as binary_file: encrypted_bytes_big = binary_file.read() - encrypted_bytes_small = self.load_bytes(small_part) + encrypted_bytes_small = self.load_bytes(small_part, folder) - encrypted_model_bytes = encrypted_bytes_small + encrypted_bytes_big - key = Security.get_model_encryption_key() + encrypted_bytes = encrypted_bytes_small + encrypted_bytes_big - model_bytes = Security.decrypt_to(encrypted_model_bytes, key) - return model_bytes + result = Security.decrypt_to(encrypted_bytes, key) + return result + + def upload_big_small_resource(self, resource, resource_name, folder, key): + big_part_name = f'{resource_name}.big' + small_part_name = f'{resource_name}.small' + + resource_encrypted = Security.encrypt_to(resource, key) + part_small_size = min(constants.SMALL_SIZE_KB * 1024, int(0.2 * len(resource_encrypted))) + part_small = resource_encrypted[:part_small_size] # slice bytes for part1 + + part_big = resource_encrypted[part_small_size:] + + self.cdn_manager.upload(constants.MODELS_FOLDER, big_part_name, part_big) + with open(path.join(folder, big_part_name), 'wb') as f: + f.write(part_big) + self.upload_file(small_part_name, part_small, constants.MODELS_FOLDER) \ No newline at end of file diff --git a/config.yaml b/config.yaml index beef3f6..10943c7 100644 --- a/config.yaml +++ b/config.yaml @@ -1,16 +1,7 @@ -cdn: - host: 'https://cdnapi.azaion.com/' - downloader_access_key: '8ynZ0rrMLL00GLBopklw' - downloader_access_secret: 'McNgEKhAJUxoa3t4WDDbCbhYPg4Qhe7FNQEKrtbk' - uploader_access_key: 'YhdHtKaq8DmvrYohetu6' - uploader_access_secret: 'nlOtjo1c4UWiMiJOjcIpR0aJFPitIhcwU6zFev7H' - bucket: 'models' - api: url: 'https://api.azaion.com' - user: 'admin@azaion.com' - pw: 'Az@1on1000Odm$n' - folder: '' + email: 'uploader@azaion.com' + password: 'Az@1on_10Upl0@der' queue: host: '188.245.120.247' diff --git a/constants.py b/constants.py index 1ec24b3..339b2b1 100644 --- a/constants.py +++ b/constants.py @@ -39,3 +39,7 @@ AI_ONNX_MODEL_FILE_SMALL = "azaion.onnx.small" AI_TENSOR_MODEL_FILE_BIG = "azaion.engine.big" AI_TENSOR_MODEL_FILE_SMALL = "azaion.engine.small" + +SMALL_SIZE_KB = 3 +CDN_CONFIG = 'cdn.yaml' +MODELS_FOLDER = 'models' diff --git a/exports.py b/exports.py index 740f9f0..57bb9b8 100644 --- a/exports.py +++ b/exports.py @@ -1,13 +1,15 @@ +import os import shutil from os import path, scandir, makedirs from pathlib import Path import random + import netron import yaml from ultralytics import YOLO import constants -from azaion_api import Api, ApiCredentials +from api_client import ApiClient, ApiCredentials from cdn_manager import CDNManager, CDNCredentials from constants import datasets_dir, processed_images_dir from security import Security @@ -24,15 +26,21 @@ def export_rknn(model_path): pass -def export_onnx(model_path): +def export_onnx(model_path, batch_size=4): model = YOLO(model_path) + onnx_path = Path(model_path).stem + '.onnx' + if path.exists(onnx_path): + os.remove(onnx_path) + model.export( format="onnx", imgsz=1280, - batch=2, + batch=batch_size, simplify=True, - nms=True) - return Path(model_path).stem + '.onnx' + nms=True, + device=0 + ) + return onnx_path def export_tensorrt(model_path): @@ -79,17 +87,25 @@ def upload_model(model_path: str, filename: str, size_small_in_kb: int=3): key = Security.get_model_encryption_key() model_encrypted = Security.encrypt_to(model_bytes, key) - part1_size = min(size_small_in_kb * 1024, int(0.9 * len(model_encrypted))) + part1_size = min(size_small_in_kb * 1024, int(0.3 * len(model_encrypted))) model_part_small = model_encrypted[:part1_size] # slice bytes for part1 model_part_big = model_encrypted[part1_size:] with open(constants.CONFIG_FILE, "r") as f: config_dict = yaml.safe_load(f) d_config = Dotdict(config_dict) - cdn_c = Dotdict(d_config.cdn) api_c = Dotdict(d_config.api) - cdn_manager = CDNManager(CDNCredentials(cdn_c.host, cdn_c.access_key, cdn_c.secret_key)) - cdn_manager.upload(cdn_c.bucket, f'{filename}.big', model_part_big) + api = ApiClient(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder)) - api = Api(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder)) - api.upload_file(f'{filename}.small', model_part_small) + yaml_bytes = api.load_bytes(constants.CDN_CONFIG, '') + data = yaml.safe_load(yaml_bytes) + creds = CDNCredentials(data["host"], + data["downloader_access_key"], + data["downloader_access_secret"], + data["uploader_access_key"], + data["uploader_access_secret"]) + + cdn_manager = CDNManager(creds) + + api.upload_file(f'{filename}.small', model_part_small, constants.MODELS_FOLDER) + cdn_manager.upload(constants.MODELS_FOLDER, f'{filename}.big', model_part_big) diff --git a/hardware_service.py b/hardware_service.py index 369b2c6..adb9cfb 100644 --- a/hardware_service.py +++ b/hardware_service.py @@ -34,7 +34,6 @@ def get_mac_address(interface="Ethernet"): def get_hardware_info(): - is_windows = os.name == 'nt' res = subprocess.check_output("ver", shell=True).decode('utf-8') if "Microsoft Windows" in res: is_windows = True diff --git a/inference/tensorrt_engine.py b/inference/tensorrt_engine.py index 31a0374..b75d303 100644 --- a/inference/tensorrt_engine.py +++ b/inference/tensorrt_engine.py @@ -8,28 +8,29 @@ import numpy as np import tensorrt as trt import pycuda.driver as cuda from inference.onnx_engine import InferenceEngine -import pycuda.autoinit # required for automatically initialize CUDA, do not remove. +# required for automatically initialize CUDA, do not remove. +import pycuda.autoinit +import pynvml +# TODO: 2. Convert onnx model with 4 batch and make sure it is working class TensorRTEngine(InferenceEngine): - def __init__(self, model_bytes: bytes, batch_size: int = 4, **kwargs): - self.batch_size = batch_size + TRT_LOGGER = trt.Logger(trt.Logger.WARNING) + def __init__(self, model_bytes: bytes, **kwargs): try: - logger = trt.Logger(trt.Logger.WARNING) + # metadata_len = struct.unpack(" int: return self.batch_size - # In tensorrt_engine.py, modify the run method: + @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: diff --git a/requirements.txt b/requirements.txt index 3ff9b32..7fad15e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,7 +9,7 @@ opencv-python matplotlib PyYAML cryptography -numpy +numpy==1.26.4 requests pyyaml boto3 diff --git a/start_inference.py b/start_inference.py index c93d022..8cd27cb 100644 --- a/start_inference.py +++ b/start_inference.py @@ -1,14 +1,23 @@ -import re - +import pycuda.driver as cuda import yaml import constants -from azaion_api import Api, ApiCredentials +from api_client import ApiClient, ApiCredentials from cdn_manager import CDNManager, CDNCredentials from inference.inference import Inference -from inference.onnx_engine import OnnxEngine from inference.tensorrt_engine import TensorRTEngine +from security import Security from utils import Dotdict + +def get_engine_filename(device_id=0): + 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 + if __name__ == "__main__": # Inference(OnnxEngine('azaion-2025-03-10.onnx', batch_size=4), # confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4') @@ -23,26 +32,19 @@ if __name__ == "__main__": # Inference(TensorRTEngine('azaion-2025-03-10_batch8.engine', batch_size=8), # confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4') - with open(constants.CONFIG_FILE, "r") as f: - config_dict = yaml.safe_load(f) - d_config = Dotdict(config_dict) - cdn_c = Dotdict(d_config.cdn) - api_c = Dotdict(d_config.api) + api_client = ApiClient() + key = Security.get_model_encryption_key() + engine_filename = TensorRTEngine.get_engine_filename() + model_bytes = api_client.load_big_small_resource(engine_filename, 'models', key) - cdn_manager = CDNManager(CDNCredentials(cdn_c.host, - cdn_c.downloader_access_key, cdn_c.downloader_access_secret, - cdn_c.uploader_access_key, cdn_c.uploader_access_secret)) - cdn_manager.download(cdn_c.bucket, constants.AI_TENSOR_MODEL_FILE_BIG) - cdn_manager.download(cdn_c.bucket, constants.AI_ONNX_MODEL_FILE_BIG) + Inference(TensorRTEngine(model_bytes), + confidence_threshold=0.5, iou_threshold=0.3).process('tests/ForAI_test.mp4') - api_client = Api(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder)) - - tensor_model_bytes = api_client.load_resource(constants.AI_TENSOR_MODEL_FILE_BIG, constants.AI_TENSOR_MODEL_FILE_SMALL) - onxx_model_bytes = api_client.load_resource(constants.AI_ONNX_MODEL_FILE_BIG, constants.AI_ONNX_MODEL_FILE_SMALL) + # cdn_manager.download(cdn_c.bucket, constants.AI_TENSOR_MODEL_FILE_BIG) + # tensor_model_bytes = api_client.load_resource(constants.AI_TENSOR_MODEL_FILE_BIG, constants.AI_TENSOR_MODEL_FILE_SMALL) # Inference(OnnxEngine(onxx_model_bytes, batch_size=4), # confidence_threshold=0.5, iou_threshold=0.3).process('tests/ForAI_test.mp4') - Inference(TensorRTEngine(tensor_model_bytes, batch_size=4), - confidence_threshold=0.5, iou_threshold=0.3).process('tests/ForAI_test.mp4') + diff --git a/train.py b/train.py index d970384..16c5023 100644 --- a/train.py +++ b/train.py @@ -7,12 +7,13 @@ from datetime import datetime from os import path, replace, listdir, makedirs, scandir from os.path import abspath from pathlib import Path +from time import sleep import yaml from ultralytics import YOLO import constants -from azaion_api import ApiCredentials, Api +from api_client import ApiCredentials, ApiClient from cdn_manager import CDNCredentials, CDNManager from constants import (processed_images_dir, processed_labels_dir, @@ -20,10 +21,11 @@ from constants import (processed_images_dir, datasets_dir, models_dir, corrupted_images_dir, corrupted_labels_dir, sample_dir) from dto.annotationClass import AnnotationClass +from inference.onnx_engine import OnnxEngine from security import Security from utils import Dotdict -from exports import export_tensorrt, upload_model +from exports import export_tensorrt, upload_model, export_onnx today_folder = f'{prefix}{datetime.now():{date_format}}' today_dataset = path.join(datasets_dir, today_folder) @@ -224,10 +226,11 @@ if __name__ == '__main__': # validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt')) # form_data_sample(500) # convert2rknn() - model_path = 'azaion.pt' - export_tensorrt(model_path) - engine_model_path = f'{Path(model_path).stem}.engine' - upload_model(engine_model_path, engine_model_path) + api_client = ApiClient() + onnx_path = export_onnx('azaion.pt') - onnx_model_path = f'{Path(model_path).stem}.onnx' - upload_model(onnx_model_path, onnx_model_path) + with open(onnx_path, 'rb') as binary_file: + onnx_bytes = binary_file.read() + + key = Security.get_model_encryption_key() + api_client.upload_big_small_resource(onnx_bytes, onnx_path, constants.MODELS_FOLDER, key)