fixed inference bugs

add DONE during inference, correct handling on C# side
This commit is contained in:
Alex Bezdieniezhnykh
2025-02-01 02:09:11 +02:00
parent e7afa96a0b
commit 739759628a
23 changed files with 324 additions and 95 deletions
+12 -7
View File
@@ -1,5 +1,3 @@
import ai_config
import msgpack
from ultralytics import YOLO
import mimetypes
import cv2
@@ -13,6 +11,7 @@ cdef class Inference:
def __init__(self, model_bytes, on_annotation):
loader = SecureModelLoader()
model_path = loader.load_model(model_bytes)
self.stop_signal = False
self.model = YOLO(<str>model_path)
self.on_annotation = on_annotation
@@ -22,19 +21,20 @@ cdef class Inference:
cdef run_inference(self, RemoteCommand cmd, int batch_size=8):
print('run inference..')
self.stop_signal = False
if self.is_video(cmd.filename):
return self._process_video(cmd, batch_size)
self._process_video(cmd, batch_size)
else:
return self._process_image(cmd)
self._process_image(cmd)
cdef _process_video(self, RemoteCommand cmd, int batch_size):
frame_count = 0
batch_frame = []
self._previous_annotation = None
v_input = cv2.VideoCapture(<str>cmd.filename)
self.ai_config = AIRecognitionConfig.from_msgpack(cmd.data)
while v_input.isOpened():
while v_input.isOpened() and not self.stop_signal:
ret, frame = v_input.read()
ms = v_input.get(cv2.CAP_PROP_POS_MSEC)
if not ret or frame is None:
@@ -51,7 +51,9 @@ cdef class Inference:
for frame, res in zip(batch_frame, results):
annotation = self.frame_to_annotation(int(frame[1]), frame[0], res.boxes)
if self.is_valid_annotation(<Annotation>annotation):
is_valid = self.is_valid_annotation(<Annotation>annotation)
print(f'Is valid annotation: {is_valid}')
if is_valid:
self._previous_annotation = annotation
self.on_annotation(cmd, annotation)
batch_frame.clear()
@@ -64,6 +66,9 @@ cdef class Inference:
annotation = self.frame_to_annotation(0, frame, res[0].boxes)
self.on_annotation(cmd, annotation)
cdef stop(self):
self.stop_signal = True
cdef frame_to_annotation(self, long time, frame, boxes: Boxes):
detections = []
for box in boxes: