rewrite to zmq push and pull patterns.

file load works, suite can start up
This commit is contained in:
Alex Bezdieniezhnykh
2025-01-23 14:37:13 +02:00
parent ce25ef38b0
commit 82b3b526a7
20 changed files with 243 additions and 208 deletions
+4 -6
View File
@@ -2,12 +2,10 @@ from ultralytics import YOLO
import mimetypes
import cv2
from ultralytics.engine.results import Boxes
from processor_command import FileCommand
from remote_command cimport RemoteCommand
from annotation cimport Detection, Annotation
cdef class Inference:
"""Handles YOLO inference using the AI model."""
def __init__(self, model_bytes, on_annotations):
self.model = YOLO(model_bytes)
self.on_annotations = on_annotations
@@ -16,13 +14,13 @@ cdef class Inference:
mime_type, _ = mimetypes.guess_type(<str>filepath)
return mime_type and mime_type.startswith("video")
cdef run_inference(self, cmd: FileCommand, int batch_size=8, int frame_skip=4):
cdef run_inference(self, RemoteCommand cmd, int batch_size=8, int frame_skip=4):
if self.is_video(cmd.filename):
return self._process_video(cmd, batch_size, frame_skip)
else:
return self._process_image(cmd)
cdef _process_video(self, cmd: FileCommand, int batch_size, int frame_skip):
cdef _process_video(self, RemoteCommand cmd, int batch_size, int frame_skip):
frame_count = 0
batch_frame = []
annotations = []
@@ -51,7 +49,7 @@ cdef class Inference:
v_input.release()
cdef _process_image(self, cmd: FileCommand):
cdef _process_image(self, RemoteCommand cmd):
frame = cv2.imread(<str>cmd.filename)
res = self.model.track(frame)
annotation = self.process_detections(0, frame, res[0].boxes)