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

fix converting model initialization

Browse Source
This commit is contained in:
Oleksandr Bezdieniezhnykh
2025-09-04 10:39:41 +03:00
parent b3665630ed
commit be77a81875
8 changed files with 97 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Azaion.Inference/inference.pyx
+55 -18
View File
@@ -11,6 +11,8 @@ from remote_command_inf cimport RemoteCommand
from annotation cimport Detection, Annotation
from ai_config cimport AIRecognitionConfig
import pynvml
from threading import Thread
from remote_command_inf cimport RemoteCommand, CommandType
cdef int tensor_gpu_index
@@ -20,7 +22,7 @@ cdef int check_tensor_gpu_index():
deviceCount = pynvml.nvmlDeviceGetCount()
if deviceCount == 0:
constants_inf.logerror('No NVIDIA GPUs found.')
constants_inf.logerror(<str>'No NVIDIA GPUs found.')
return -1
for i in range(deviceCount):
@@ -28,10 +30,10 @@ cdef int check_tensor_gpu_index():
major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
if major > 6 or (major == 6 and minor >= 1):
constants_inf.log('found NVIDIA GPU!')
constants_inf.log(<str>'found NVIDIA GPU!')
return i
constants_inf.logerror('NVIDIA GPU doesnt support TensorRT!')
constants_inf.logerror(<str>'NVIDIA GPU doesnt support TensorRT!')
return -1
except pynvml.NVMLError:
@@ -40,7 +42,7 @@ cdef int check_tensor_gpu_index():
try:
pynvml.nvmlShutdown()
except:
constants_inf.logerror('Failed to shutdown pynvml cause probably no NVIDIA GPU')
constants_inf.logerror(<str>'Failed to shutdown pynvml cause probably no NVIDIA GPU')
pass
tensor_gpu_index = check_tensor_gpu_index()
@@ -51,9 +53,9 @@ else:
cdef class Inference:
def __init__(self, loader_client, on_annotation):
def __init__(self, loader_client, remote_handler):
self.loader_client = loader_client
self.on_annotation = on_annotation
self.remote_handler = remote_handler
self.stop_signal = False
self.model_input = None
self.model_width = 0
@@ -61,8 +63,10 @@ cdef class Inference:
self.engine = None
self.is_building_engine = False
self.ai_availability_status = AIAvailabilityStatus()
self._converted_model_bytes = None
self.init_ai()
cdef bytes get_onnx_engine_bytes(self):
models_dir = constants_inf.MODELS_FOLDER
self.ai_availability_status.set_status(AIAvailabilityEnum.DOWNLOADING)
@@ -71,15 +75,43 @@ cdef class Inference:
raise Exception(res.err)
return res.data
cdef convert_and_upload_model(self, bytes onnx_engine_bytes, str engine_filename):
try:
self.ai_availability_status.set_status(AIAvailabilityEnum.CONVERTING)
models_dir = constants_inf.MODELS_FOLDER
model_bytes = TensorRTEngine.convert_from_onnx(onnx_engine_bytes)
self.ai_availability_status.set_status(AIAvailabilityEnum.UPLOADING)
res = self.loader_client.upload_big_small_resource(model_bytes, engine_filename, models_dir)
if res.err is not None:
self.ai_availability_status.set_status(AIAvailabilityEnum.WARNING, <str>f"Failed to upload converted model: {res.err}")
self._converted_model_bytes = model_bytes
except Exception as e:
self.ai_availability_status.set_status(AIAvailabilityEnum.ERROR, <str> str(e))
self._converted_model_bytes = None
finally:
self.is_building_engine = False
cdef init_ai(self):
constants_inf.log(<str> 'init AI...')
try:
while self.is_building_engine:
time.sleep(1)
if self.engine is not None:
return
if self.is_building_engine:
return
if self._converted_model_bytes is not None:
try:
self.engine = TensorRTEngine(self._converted_model_bytes)
self.ai_availability_status.set_status(AIAvailabilityEnum.ENABLED)
self.model_height, self.model_width = self.engine.get_input_shape()
except Exception as e:
self.ai_availability_status.set_status(AIAvailabilityEnum.ERROR, <str> str(e))
finally:
self._converted_model_bytes = None # Consume the bytes
return
self.is_building_engine = True
models_dir = constants_inf.MODELS_FOLDER
if tensor_gpu_index > -1:
try:
@@ -93,15 +125,12 @@ cdef class Inference:
except Exception as e:
self.ai_availability_status.set_status(AIAvailabilityEnum.WARNING, <str>str(e))
onnx_engine_bytes = self.get_onnx_engine_bytes()
self.ai_availability_status.set_status(AIAvailabilityEnum.CONVERTING)
model_bytes = TensorRTEngine.convert_from_onnx(onnx_engine_bytes)
self.engine = TensorRTEngine(model_bytes)
self.ai_availability_status.set_status(AIAvailabilityEnum.UPLOADING)
res = self.loader_client.upload_big_small_resource(model_bytes, <str> engine_filename, models_dir)
if res.err is not None:
self.ai_availability_status.set_status(AIAvailabilityEnum.ERROR, res.err)
else:
self.ai_availability_status.set_status(AIAvailabilityEnum.ENABLED)
self.is_building_engine = True
thread = Thread(target=self.convert_and_upload_model, args=(onnx_engine_bytes, engine_filename))
thread.daemon = True
thread.start()
return
else:
self.engine = OnnxEngine(<bytes>self.get_onnx_engine_bytes())
self.is_building_engine = False
@@ -200,6 +229,11 @@ cdef class Inference:
self.stop_signal = False
self.init_ai()
if self.engine is None:
constants_inf.log(<str> "AI engine not available. Conversion may be in progress. Skipping inference.")
response = RemoteCommand(CommandType.AI_AVAILABILITY_RESULT, self.ai_availability_status.serialize())
self.remote_handler.send(cmd.client_id, response.serialize())
return
for m in ai_config.paths:
if self.is_video(m):
@@ -258,6 +292,9 @@ cdef class Inference:
batch_timestamps.clear()
v_input.release()
cdef on_annotation(self, RemoteCommand cmd, Annotation annotation):
cdef RemoteCommand response = RemoteCommand(CommandType.INFERENCE_DATA, annotation.serialize())
self.remote_handler.send(cmd.client_id, response.serialize())
cdef _process_images(self, RemoteCommand cmd, AIRecognitionConfig ai_config, list[str] image_paths):
cdef list frame_data