mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 12:16:30 +00:00
fix bug with annotation result gradient stops
add tensorrt engine
This commit is contained in:
Alex Bezdieniezhnykh
@@ -1,40 +1,40 @@
|
||||
import json
|
||||
import mimetypes
|
||||
import subprocess
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import onnxruntime as onnx
|
||||
|
||||
cimport constants
|
||||
from remote_command cimport RemoteCommand
|
||||
from annotation cimport Detection, Annotation
|
||||
from ai_config cimport AIRecognitionConfig
|
||||
from inference_engine cimport OnnxEngine, TensorRTEngine
|
||||
from hardware_service cimport HardwareService
|
||||
|
||||
cdef class Inference:
|
||||
def __init__(self, api_client, on_annotation):
|
||||
self.api_client = api_client
|
||||
self.on_annotation = on_annotation
|
||||
self.stop_signal = False
|
||||
self.session = None
|
||||
self.model_input = None
|
||||
self.model_width = 0
|
||||
self.model_height = 0
|
||||
self.engine = None
|
||||
self.class_names = None
|
||||
|
||||
def init_ai(self):
|
||||
model_bytes = self.api_client.load_ai_model()
|
||||
self.session = onnx.InferenceSession(
|
||||
model_bytes, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
)
|
||||
model_inputs = self.session.get_inputs()
|
||||
self.model_input = model_inputs[0].name
|
||||
input_shape = model_inputs[0].shape
|
||||
self.model_width = input_shape[2]
|
||||
self.model_height = input_shape[3]
|
||||
print(f'AI detection model input: {self.model_input} ({self.model_width}, {self.model_height})')
|
||||
model_meta = self.session.get_modelmeta()
|
||||
print("Metadata:", model_meta.custom_metadata_map)
|
||||
self.class_names = eval(model_meta.custom_metadata_map["names"])
|
||||
if self.engine is not None:
|
||||
return
|
||||
|
||||
is_nvidia = HardwareService.has_nvidia_gpu()
|
||||
if is_nvidia:
|
||||
model_bytes = self.api_client.load_ai_model(is_tensor=True)
|
||||
self.engine = TensorRTEngine(model_bytes, batch_size=4)
|
||||
else:
|
||||
model_bytes = self.api_client.load_ai_model()
|
||||
self.engine = OnnxEngine(model_bytes, batch_size=4)
|
||||
|
||||
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,
|
||||
@@ -47,33 +47,37 @@ 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()
|
||||
for det in output[0][ann_index]:
|
||||
if det[4] == 0: # if confidence is 0 then valid points are over.
|
||||
break
|
||||
x1 = det[0] / self.model_width
|
||||
y1 = det[1] / self.model_height
|
||||
x2 = det[2] / self.model_width
|
||||
y2 = det[3] / self.model_height
|
||||
conf = round(det[4], 2)
|
||||
class_id = int(det[5])
|
||||
|
||||
for ann_index in range(len(output[0])):
|
||||
detections.clear()
|
||||
for det in output[0][ann_index]:
|
||||
if det[4] == 0: # if confidence is 0 then valid points are over.
|
||||
break
|
||||
x1 = det[0] / self.model_width
|
||||
y1 = det[1] / self.model_height
|
||||
x2 = det[2] / self.model_width
|
||||
y2 = det[3] / self.model_height
|
||||
conf = round(det[4], 2)
|
||||
class_id = int(det[5])
|
||||
|
||||
x = (x1 + x2) / 2
|
||||
y = (y1 + y2) / 2
|
||||
w = x2 - x1
|
||||
h = y2 - y1
|
||||
if conf >= ai_config.probability_threshold:
|
||||
detections.append(Detection(x, y, w, h, class_id, conf))
|
||||
filtered_detections = self.remove_overlapping_detections(detections)
|
||||
results.append(filtered_detections)
|
||||
return results
|
||||
x = (x1 + x2) / 2
|
||||
y = (y1 + y2) / 2
|
||||
w = x2 - x1
|
||||
h = y2 - y1
|
||||
if conf >= ai_config.probability_threshold:
|
||||
detections.append(Detection(x, y, w, h, class_id, conf))
|
||||
filtered_detections = self.remove_overlapping_detections(detections)
|
||||
results.append(filtered_detections)
|
||||
return results
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to postprocess: {str(e)}")
|
||||
|
||||
cdef remove_overlapping_detections(self, list[Detection] detections):
|
||||
cdef Detection det1, det2
|
||||
@@ -121,8 +125,7 @@ cdef class Inference:
|
||||
raise Exception('ai recognition config is empty')
|
||||
|
||||
self.stop_signal = False
|
||||
if self.session is None:
|
||||
self.init_ai()
|
||||
self.init_ai()
|
||||
|
||||
print(ai_config.paths)
|
||||
for m in ai_config.paths:
|
||||
@@ -160,7 +163,9 @@ cdef class Inference:
|
||||
|
||||
if len(batch_frames) == ai_config.model_batch_size:
|
||||
input_blob = self.preprocess(batch_frames)
|
||||
outputs = self.session.run(None, {self.model_input: input_blob})
|
||||
|
||||
outputs = self.engine.run(input_blob)
|
||||
|
||||
list_detections = self.postprocess(outputs, ai_config)
|
||||
for i in range(len(list_detections)):
|
||||
detections = list_detections[i]
|
||||
@@ -189,7 +194,9 @@ cdef class Inference:
|
||||
timestamps.append(0)
|
||||
|
||||
input_blob = self.preprocess(frames)
|
||||
outputs = self.session.run(None, {self.model_input: input_blob})
|
||||
|
||||
outputs = self.engine.run(input_blob)
|
||||
|
||||
list_detections = self.postprocess(outputs, ai_config)
|
||||
for i in range(len(list_detections)):
|
||||
detections = list_detections[i]
|
||||
@@ -199,6 +206,7 @@ cdef class Inference:
|
||||
print(annotation.to_str(self.class_names))
|
||||
self.on_annotation(cmd, annotation)
|
||||
|
||||
|
||||
cdef stop(self):
|
||||
self.stop_signal = True
|
||||
|
||||
|
||||
Reference in New Issue
Block a user