fix bug with annotation result gradient stops

add tensorrt engine

Azaion.Inference/inference_engine.pyx

@@ -0,0 +1,140 @@
+import json
+import struct
+from typing import List, Tuple
+import numpy as np
+import onnxruntime as onnx
+import tensorrt as trt
+import pycuda.driver as cuda
+import pycuda.autoinit # required for automatically initialize CUDA, do not remove.
+
+
+cdef class InferenceEngine:
+    def __init__(self, model_bytes: bytes, batch_size: int = 1, **kwargs):
+        self.batch_size = batch_size
+
+    cdef tuple get_input_shape(self):
+        raise NotImplementedError("Subclass must implement get_input_shape")
+
+    cdef int get_batch_size(self):
+        return self.batch_size
+
+    cpdef run(self, input_data):
+        raise NotImplementedError("Subclass must implement run")
+
+    cdef get_class_names(self):
+        raise NotImplementedError("Subclass must implement get_class_names")
+
+
+cdef class OnnxEngine(InferenceEngine):
+    def __init__(self, model_bytes: bytes, batch_size: int = 1, **kwargs):
+        super().__init__(model_bytes, batch_size)
+        self.batch_size = batch_size
+        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]
+        print(f'AI detection model input: {self.model_inputs} {self.input_shape}')
+        model_meta = self.session.get_modelmeta()
+        print("Metadata:", model_meta.custom_metadata_map)
+        self.class_names = eval(model_meta.custom_metadata_map["names"])
+
+    cdef tuple get_input_shape(self):
+        shape = self.input_shape
+        return shape[2], shape[3]
+
+    cdef int get_batch_size(self):
+        return self.batch_size
+
+    cdef get_class_names(self):
+        return self.class_names
+
+    cpdef run(self, input_data):
+        return self.session.run(None, {self.input_name: input_data})
+
+
+cdef class TensorRTEngine(InferenceEngine):
+    def __init__(self, model_bytes: bytes, batch_size: int = 4, **kwargs):
+        super().__init__(model_bytes, batch_size)
+        self.batch_size = batch_size
+        print('Enter init TensorRT')
+        try:
+            logger = trt.Logger(trt.Logger.WARNING)
+
+            metadata_len = struct.unpack("<I", model_bytes[:4])[0]
+            try:
+                metadata = json.loads(model_bytes[4:4 + metadata_len])
+                print(f"Model metadata: {json.dumps(metadata, indent=2)}")
+                string_dict = metadata['names']
+                self.class_names = {int(k): v for k, v in string_dict.items()}
+            except json.JSONDecodeError:
+                print(f"Failed to parse metadata")
+                return
+            engine_data = model_bytes[4 + metadata_len:]
+
+
+            runtime = trt.Runtime(logger)
+            engine = runtime.deserialize_cuda_engine(engine_data)
+
+            if engine is None:
+                raise RuntimeError(f"Failed to load TensorRT engine from bytes")
+
+            self.context = engine.create_execution_context()
+            # input
+            self.input_name = engine.get_tensor_name(0)
+            engine_input_shape = 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 = engine.get_tensor_name(1)
+            engine_output_shape = tuple(engine.get_tensor_shape(self.output_name))
+            self.output_shape = [
+                batch_size if self.input_shape[0] == -1 else self.input_shape[0],
+                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)}")
+
+    cdef tuple get_input_shape(self):
+        return self.input_shape[2], self.input_shape[3]
+
+    cdef int get_batch_size(self):
+        return self.batch_size
+
+    cdef get_class_names(self):
+        return self.class_names
+
+    cpdef run(self, input_data):
+        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)}")