azaion/annotations
1
0
Fork 0
mirror of https://github.com/azaion/annotations.git synced 2026-04-22 12:46:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

autoconvert tensor rt engine from onnx to specific CUDA gpu

Browse Source
This commit is contained in:
Alex Bezdieniezhnykh
2025-04-24 16:30:21 +03:00
parent e798af470b
commit e9a44e368d
14 changed files with 25 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Azaion.Inference/annotation.pxd
-2
View File
@@ -14,5 +14,3 @@ cdef class Annotation:
cdef format_time(self, ms)
cdef bytes serialize(self)
cdef to_str(self, class_names)
Azaion.Inference/annotation.pyx
+2 -2
View File
@@ -32,12 +32,12 @@ cdef class Annotation:
d.annotation_name = self.name
self.image = b''
cdef to_str(self, class_names):
def __str__(self):
if not self.detections:
return f"{self.name}: No detections"
detections_str = ", ".join(
f"{class_names[d.cls]} {d.confidence * 100:.1f}% ({d.x:.2f}, {d.y:.2f})"
f"class: {d.cls} {d.confidence * 100:.1f}% ({d.x:.2f}, {d.y:.2f})"
for d in self.detections
)
return f"{self.name}: {detections_str}"
Azaion.Inference/api_client.pyx
+1 -1
View File
@@ -79,7 +79,7 @@ cdef class ApiClient:
try:
r = requests.post(url, headers=headers, files=files, allow_redirects=True)
r.raise_for_status()
print(f"Upload success: {r.status_code}")
print(f"Uploaded {filename} to {constants.API_URL}/{folder} successfully: {r.status_code}.")
except Exception as e:
print(f"Upload fail: {e}")
Azaion.Inference/inference.pxd
-1
View File
@@ -10,7 +10,6 @@ cdef class Inference:
cdef object on_annotation
cdef Annotation _previous_annotation
cdef AIRecognitionConfig ai_config
cdef object class_names
cdef bint stop_signal
cdef str model_input
Azaion.Inference/inference.pyx
+5 -9
View File
@@ -24,7 +24,6 @@ cdef class Inference:
self.model_width = 0
self.model_height = 0
self.engine = None
self.class_names = None
self.is_building_engine = False
cdef build_tensor_engine(self):
@@ -37,9 +36,11 @@ cdef class Inference:
models_dir = constants.MODELS_FOLDER
if not os.path.exists(os.path.join(<str> models_dir, f'{engine_filename}.big')):
self.is_building_engine = True
time.sleep(5) # prevent simultaneously loading dll and models
onnx_model = self.api_client.load_big_small_resource(constants.AI_ONNX_MODEL_FILE, models_dir, key)
model_bytes = TensorRTEngine.convert_from_onnx(onnx_model)
self.api_client.upload_big_small_resource(model_bytes, <str> engine_filename, models_dir, key)
print('uploaded ')
self.is_building_engine = False
@@ -62,7 +63,6 @@ cdef class Inference:
self.engine = OnnxEngine(model_bytes)
self.model_height, self.model_width = self.engine.get_input_shape()
self.class_names = self.engine.get_class_names()
cdef preprocess(self, frames):
blobs = [cv2.dnn.blobFromImage(frame,
@@ -75,13 +75,11 @@ cdef class Inference:
return np.vstack(blobs)
cdef postprocess(self, output, ai_config):
print('enter postprocess')
cdef list[Detection] detections = []
cdef int ann_index
cdef float x1, y1, x2, y2, conf, cx, cy, w, h
cdef int class_id
cdef list[list[Detection]] results = []
print('start try: code')
try:
for ann_index in range(len(output[0])):
detections.clear()
@@ -163,7 +161,7 @@ cdef class Inference:
images.append(m)
# images first, it's faster
if len(images) > 0:
for chunk in self.split_list_extend(images, self.engine.get_input_shape()):
for chunk in self.split_list_extend(images, self.engine.get_batch_size()):
print(f'run inference on {" ".join(chunk)}...')
self._process_images(cmd, ai_config, chunk)
if len(videos) > 0:
@@ -189,7 +187,7 @@ cdef class Inference:
batch_frames.append(frame)
batch_timestamps.append(int(v_input.get(cv2.CAP_PROP_POS_MSEC)))
if len(batch_frames) == self.engine.get_input_shape():
if len(batch_frames) == self.engine.get_batch_size():
input_blob = self.preprocess(batch_frames)
outputs = self.engine.run(input_blob)
@@ -203,10 +201,9 @@ cdef class Inference:
annotation.image = image.tobytes()
self._previous_annotation = annotation
print(annotation.to_str(self.class_names))
print(annotation)
self.on_annotation(cmd, annotation)
batch_frames.clear()
batch_timestamps.clear()
v_input.release()
@@ -231,7 +228,6 @@ cdef class Inference:
annotation = Annotation(image_paths[i], timestamps[i], detections)
_, image = cv2.imencode('.jpg', frames[i])
annotation.image = image.tobytes()
print(annotation.to_str(self.class_names))
self.on_annotation(cmd, annotation)
Azaion.Inference/inference_engine.pxd
-3
View File
@@ -6,7 +6,6 @@ cdef class InferenceEngine:
cdef public int batch_size
cdef tuple get_input_shape(self)
cdef int get_batch_size(self)
cdef get_class_names(self)
cpdef run(self, input_data)
cdef class OnnxEngine(InferenceEngine):
@@ -14,7 +13,6 @@ cdef class OnnxEngine(InferenceEngine):
cdef list model_inputs
cdef str input_name
cdef object input_shape
cdef object class_names
cdef class TensorRTEngine(InferenceEngine):
cdef object stream
@@ -26,7 +24,6 @@ cdef class TensorRTEngine(InferenceEngine):
cdef object input_shape
cdef object output_shape
cdef object h_output
cdef object class_names
@staticmethod
cdef bytes convert_from_onnx(bytes onnx_model)
Azaion.Inference/inference_engine.pyx
+6 -12
View File
@@ -22,9 +22,6 @@ cdef class InferenceEngine:
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):
@@ -39,7 +36,6 @@ cdef class OnnxEngine(InferenceEngine):
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
@@ -48,9 +44,6 @@ cdef class OnnxEngine(InferenceEngine):
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})
@@ -69,10 +62,14 @@ cdef class TensorRTEngine(InferenceEngine):
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)
self.batch_size = self.input_shape[0] if self.input_shape[0] != -1 else batch_size
if engine_input_shape[0] != -1:
self.batch_size = engine_input_shape[0]
else:
self.batch_size = batch_size
self.input_shape = [
self.batch_size,
@@ -154,7 +151,7 @@ cdef class TensorRTEngine(InferenceEngine):
if plan is None:
print('Conversion failed.')
return None
print('conversion done!')
return bytes(plan)
cdef tuple get_input_shape(self):
@@ -163,9 +160,6 @@ cdef class TensorRTEngine(InferenceEngine):
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)
Azaion.Inference/remote_command_handler.pyx
+1 -1
View File
@@ -76,7 +76,7 @@ cdef class RemoteCommandHandler:
with self._context.socket(zmq.DEALER) as socket:
socket.connect("inproc://backend")
socket.send_multipart([client_id, data])
constants.log(<str>f'Sent {len(data)} bytes.', client_id)
# constants.log(<str>f'Sent {len(data)} bytes.', client_id)
cdef stop(self):
self._shutdown_event.set()