mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 18:36:32 +00:00
rewrite to zmq push and pull patterns.
file load works, suite can start up
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user