From 06a23525a6a6526515b9e76f2b6297b91a967ca4 Mon Sep 17 00:00:00 2001 From: zxsanny Date: Fri, 28 Mar 2025 14:50:43 +0200 Subject: [PATCH] fix tensor rt engine --- README.md | 5 ++- azaion_api.py | 59 ++++++++++++++++++++++++++++- classes.json | 8 ++-- constants.py | 8 +++- exports/export.py => exports.py | 33 ++++++++++++++--- hardware_service.py | 66 +++++++++++++++++++++++++++++++++ inference/__init__.py | 0 inference/dto.py | 1 + inference/inference.py | 5 +-- inference/onnx_engine.py | 12 ++++-- inference/start_inference.py | 20 ---------- inference/tensorrt_engine.py | 34 +++++++++-------- requirements.txt | 7 +--- security.py | 20 +++++++--- start_inference.py | 49 ++++++++++++++++++++++++ train.py | 39 +++++-------------- 16 files changed, 272 insertions(+), 94 deletions(-) rename exports/export.py => exports.py (61%) create mode 100644 hardware_service.py create mode 100644 inference/__init__.py delete mode 100644 inference/start_inference.py create mode 100644 start_inference.py diff --git a/README.md b/README.md index 2f82476..c3e006e 100644 --- a/README.md +++ b/README.md @@ -26,12 +26,15 @@ Linux python -m venv env env\Scripts\activate pip install -r requirements.txt + + pip uninstall -y torch + pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126 ``` **3. Fix possible problems** * cv2.error: OpenCV(4.10.0) ...\window.cpp:1301: error: (-2:Unspecified error) ``` - pip uninstall opencv-python + pip uninstall -y opencv-python pip install opencv-python ``` * fbgemm.dll error (Windows specific) diff --git a/azaion_api.py b/azaion_api.py index 2588b0f..14b4f0c 100644 --- a/azaion_api.py +++ b/azaion_api.py @@ -1,6 +1,13 @@ import io +import json +from http import HTTPStatus + import requests +import constants +from hardware_service import get_hardware_info +from security import Security + class ApiCredentials: def __init__(self, url, email, password, folder): @@ -34,4 +41,54 @@ class Api: r.raise_for_status() print(f"Upload {len(file_bytes)} bytes ({filename}) to {self.credentials.url}. Result: {r.status_code}") except Exception as e: - print(f"Upload fail: {e}") \ No newline at end of file + print(f"Upload fail: {e}") + + def load_bytes(self, filename, folder = None): + folder = folder or self.credentials.folder + + hardware = get_hardware_info() + + if self.token is None: + self.login() + url = f"{self.credentials.url}/resources/get/{folder}" + headers = { + "Authorization": f"Bearer {self.token}", + "Content-Type": "application/json" + } + + payload = json.dumps( + { + "password": self.credentials.password, + "hardware": hardware.to_json_object(), + "fileName": filename + }, indent=4) + response = requests.post(url, data=payload, headers=headers, stream=True) + if response.status_code == HTTPStatus.UNAUTHORIZED or response.status_code == HTTPStatus.FORBIDDEN: + self.login() + headers = { + "Authorization": f"Bearer {self.token}", + "Content-Type": "application/json" + } + response = requests.post(url, data=payload, headers=headers, stream=True) + + if response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR: + print('500!') + + hw_hash = Security.get_hw_hash(hardware) + key = Security.get_api_encryption_key(self.credentials, hw_hash) + + resp_bytes = response.raw.read() + data = Security.decrypt_to(resp_bytes, key) + print(f'Downloaded file: {filename}, {len(data)} bytes') + return data + + def load_resource(self, big_part, small_part): + with open(big_part, 'rb') as binary_file: + encrypted_bytes_big = binary_file.read() + encrypted_bytes_small = self.load_bytes(small_part) + + encrypted_model_bytes = encrypted_bytes_small + encrypted_bytes_big + key = Security.get_model_encryption_key() + + model_bytes = Security.decrypt_to(encrypted_model_bytes, key) + return model_bytes diff --git a/classes.json b/classes.json index a83dd9e..eff5acc 100644 --- a/classes.json +++ b/classes.json @@ -5,12 +5,14 @@ { "Id": 3, "Name": "Artillery", "Color": "#80FFFF00" }, { "Id": 4, "Name": "Shadow", "Color": "#80FF00FF" }, { "Id": 5, "Name": "Trenches", "Color": "#8000FFFF" }, - { "Id": 6, "Name": "Military-men", "Color": "#80000000" }, + { "Id": 6, "Name": "Military-men", "Color": "#80188021" }, { "Id": 7, "Name": "Tyre-tracks", "Color": "#80800000" }, { "Id": 8, "Name": "Additional-armored-tank", "Color": "#80008000" }, { "Id": 9, "Name": "Smoke", "Color": "#80000080" }, { "Id": 10, "Name": "Plane", "Color": "#80000080" }, { "Id": 11, "Name": "Moto", "Color": "#80808000" }, { "Id": 12, "Name": "Camouflage-net", "Color": "#80800080" }, - { "Id": 13, "Name": "Camouflage-branches", "Color": "#80008080" } - ] + { "Id": 13, "Name": "Camouflage-branches", "Color": "#80008080" }, + { "Id": 14, "Name": "Roof", "Color": "#80008080" }, + { "Id": 15, "Name": "Building", "Color": "#80008080" } + ] diff --git a/constants.py b/constants.py index 5beff4a..1ec24b3 100644 --- a/constants.py +++ b/constants.py @@ -1,5 +1,4 @@ from os import path -from dto.annotationClass import AnnotationClass azaion = '/azaion' prefix = 'azaion-' @@ -26,7 +25,6 @@ datasets_dir = path.join(azaion, 'datasets') models_dir = path.join(azaion, 'models') -annotation_classes = AnnotationClass.read_json() date_format = '%Y-%m-%d' checkpoint_file = 'checkpoint.txt' checkpoint_date_format = '%Y-%m-%d %H:%M:%S' @@ -35,3 +33,9 @@ CONFIG_FILE = 'config.yaml' ANNOTATIONS_QUEUE = 'azaion-annotations' ANNOTATIONS_CONFIRMED_QUEUE = 'azaion-annotations-confirm' OFFSET_FILE = 'offset.yaml' + +AI_ONNX_MODEL_FILE_BIG = "azaion.onnx.big" +AI_ONNX_MODEL_FILE_SMALL = "azaion.onnx.small" + +AI_TENSOR_MODEL_FILE_BIG = "azaion.engine.big" +AI_TENSOR_MODEL_FILE_SMALL = "azaion.engine.small" diff --git a/exports/export.py b/exports.py similarity index 61% rename from exports/export.py rename to exports.py index 88ae2ae..740f9f0 100644 --- a/exports/export.py +++ b/exports.py @@ -2,12 +2,16 @@ 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 cdn_manager import CDNManager, CDNCredentials from constants import datasets_dir, processed_images_dir +from security import Security +from utils import Dotdict def export_rknn(model_path): @@ -40,6 +44,7 @@ def export_tensorrt(model_path): nms=True ) + def form_data_sample(destination_path, size=500, write_txt_log=False): images = [] with scandir(processed_images_dir) as imd: @@ -62,11 +67,29 @@ def form_data_sample(destination_path, size=500, write_txt_log=False): with open(path.join(destination_path, 'azaion_subset.txt'), 'w', encoding='utf-8') as f: f.writelines([f'{line}\n' for line in lines]) + def show_model(model: str = None): netron.start(model) -if __name__ == '__main__': - export_tensorrt('azaion-2025-03-10.pt') - # export_rknn('azaion-2025-03-10.pt') - # export_onnx('azaion-2025-03-10.pt') \ No newline at end of file +def upload_model(model_path: str, filename: str, size_small_in_kb: int=3): + with open(model_path, 'rb') as f_in: + model_bytes = f_in.read() + + 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))) + 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 = Api(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder)) + api.upload_file(f'{filename}.small', model_part_small) diff --git a/hardware_service.py b/hardware_service.py new file mode 100644 index 0000000..369b2c6 --- /dev/null +++ b/hardware_service.py @@ -0,0 +1,66 @@ +import os +import subprocess + +import psutil + + +class HardwareInfo: + def __init__(self, cpu, gpu, memory, mac_address): + self.cpu = cpu + self.gpu = gpu + self.memory = memory + self.mac_address = mac_address + + def to_json_object(self): + return { + "CPU": self.cpu, + "GPU": self.gpu, + "MacAddress": self.mac_address, + "Memory": self.memory + } + + def __str__(self): + return f'CPU: {self.cpu}. GPU: {self.gpu}. Memory: {self.memory}. MAC Address: {self.mac_address}' + + +def get_mac_address(interface="Ethernet"): + addresses = psutil.net_if_addrs() + for interface_name, interface_info in addresses.items(): + if interface_name == interface: + for addr in interface_info: + if addr.family == psutil.AF_LINK: + return addr.address.replace('-', '') + return None + + +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 + else: + is_windows = False + + if is_windows: + os_command = ( + "powershell -Command \"" + "Get-CimInstance -ClassName Win32_Processor | Select-Object -ExpandProperty Name | Write-Output; " + "Get-CimInstance -ClassName Win32_VideoController | Select-Object -ExpandProperty Name | Write-Output; " + "Get-CimInstance -ClassName Win32_OperatingSystem | Select-Object -ExpandProperty TotalVisibleMemorySize | Write-Output" + "\"" + ) + else: + os_command = ( + "/bin/bash -c \" lscpu | grep 'Model name:' | cut -d':' -f2 && " + "lspci | grep VGA | cut -d':' -f3 && " + "free -g | grep Mem: | awk '{print $2}' && \"" + ) + result = subprocess.check_output(os_command, shell=True).decode('utf-8') + lines = [line.strip() for line in result.splitlines() if line.strip()] + + cpu = lines[0].replace("Name=", "").replace(" ", " ") + gpu = lines[1].replace("Name=", "").replace(" ", " ") + memory = lines[2].replace("TotalVisibleMemorySize=", "").replace(" ", " ") + mac_address = get_mac_address() + + return HardwareInfo(cpu, gpu, memory, mac_address) \ No newline at end of file diff --git a/inference/__init__.py b/inference/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/inference/dto.py b/inference/dto.py index 1f7cae9..cd1b2ab 100644 --- a/inference/dto.py +++ b/inference/dto.py @@ -33,6 +33,7 @@ class WeatherMode(Enum): Wint = 20 Night = 40 + class AnnotationClass: def __init__(self, id, name, color): self.id = id diff --git a/inference/inference.py b/inference/inference.py index 5405051..7b0c095 100644 --- a/inference/inference.py +++ b/inference/inference.py @@ -1,8 +1,7 @@ import cv2 import numpy as np - -from onnx_engine import InferenceEngine -from dto import AnnotationClass, Annotation, Detection +from inference.dto import Annotation, Detection, AnnotationClass +from inference.onnx_engine import InferenceEngine class Inference: diff --git a/inference/onnx_engine.py b/inference/onnx_engine.py index 3af914f..a4e9ba6 100644 --- a/inference/onnx_engine.py +++ b/inference/onnx_engine.py @@ -22,15 +22,19 @@ class InferenceEngine(abc.ABC): pass - class OnnxEngine(InferenceEngine): - def __init__(self, model_path: str, batch_size: int = 1, **kwargs): - self.model_path = model_path + def __init__(self, model_bytes, batch_size: int = 1, **kwargs): self.batch_size = batch_size - self.session = onnx.InferenceSession(model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) + self.session = onnx.InferenceSession(model_bytes, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) self.model_inputs = self.session.get_inputs() self.input_name = self.model_inputs[0].name self.input_shape = self.model_inputs[0].shape + if self.input_shape[0] != -1: + self.batch_size = self.input_shape[0] + model_meta = self.session.get_modelmeta() + print("Metadata:", model_meta.custom_metadata_map) + self.class_names = eval(model_meta.custom_metadata_map["names"]) + pass def get_input_shape(self) -> Tuple[int, int]: shape = self.input_shape diff --git a/inference/start_inference.py b/inference/start_inference.py deleted file mode 100644 index 7d1f11d..0000000 --- a/inference/start_inference.py +++ /dev/null @@ -1,20 +0,0 @@ -from onnx_engine import OnnxEngine -from tensorrt_engine import TensorRTEngine -from inference import Inference - -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') - # detection for the first 200sec of video: - # onnxInference: 81 sec, 6.3Gb VRAM - # tensorrt: 54 sec, 3.7Gb VRAM - - # Inference(TensorRTEngine('azaion-2025-03-10_int8.engine', batch_size=16), - # confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4') - # INT8 for 200sec: 54 sec 3.7Gb - - # Inference(TensorRTEngine('azaion-2025-03-10_batch8.engine', batch_size=8), - # confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4') - - Inference(TensorRTEngine('azaion-2025-03-10-half_batch4.engine', batch_size=4), - confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4') \ No newline at end of file diff --git a/inference/tensorrt_engine.py b/inference/tensorrt_engine.py index 0dfff06..31a0374 100644 --- a/inference/tensorrt_engine.py +++ b/inference/tensorrt_engine.py @@ -1,46 +1,48 @@ +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 import pycuda.autoinit # required for automatically initialize CUDA, do not remove. -from onnx_engine import InferenceEngine - class TensorRTEngine(InferenceEngine): - def __init__(self, model_path: str, batch_size: int = 4, **kwargs): - self.model_path = model_path + def __init__(self, model_bytes: bytes, batch_size: int = 4, **kwargs): self.batch_size = batch_size try: logger = trt.Logger(trt.Logger.WARNING) - with open(model_path, 'rb') as f: - metadata_len = int.from_bytes(f.read(4), byteorder='little', signed=True) - metadata_bytes = f.read(metadata_len) - try: - self.metadata = json.loads(metadata_bytes) - print(f"Model metadata: {json.dumps(self.metadata, indent=2)}") - except json.JSONDecodeError: - print(f"Failed to parse metadata: {metadata_bytes}") - self.metadata = {} - engine_data = f.read() + metadata_len = struct.unpack("