mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 12:56:30 +00:00
fixed inference bugs
add DONE during inference, correct handling on C# side
This commit is contained in:
+12
-7
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user