annotations/Azaion.Inference/main.pyx

import traceback
from queue import Queue
cimport constants
import msgpack

from api_client cimport ApiClient
from annotation cimport Annotation
from inference cimport Inference
from remote_command cimport RemoteCommand, CommandType
from remote_command_handler cimport RemoteCommandHandler
from user cimport User

cdef class ParsedArguments:
    cdef str email, password, folder;

    def __init__(self, str email, str password, str folder):
        self.email = email
        self.password = password
        self.folder = folder

cdef class CommandProcessor:
    cdef ApiClient api_client
    cdef RemoteCommandHandler remote_handler
    cdef object command_queue
    cdef bint running
    cdef Inference inference

    def __init__(self, args: ParsedArguments):
        self.api_client = ApiClient(args.email, args.password, args.folder)
        self.remote_handler = RemoteCommandHandler(self.on_command)
        self.command_queue = Queue(maxsize=constants.QUEUE_MAXSIZE)
        self.remote_handler.start()
        self.running = True
        model = self.api_client.load_ai_model()
        self.inference = Inference(model, self.on_annotation)

    def start(self):
        while self.running:
            try:
                command = self.command_queue.get()
                self.inference.run_inference(command)
                self.remote_handler.send(command.client_id, <bytes>'DONE'.encode('utf-8'))
            except Exception as e:
                traceback.print_exc()

    cdef on_command(self, RemoteCommand command):
        try:
            if command.command_type == CommandType.GET_USER:
                self.get_user(command, self.api_client.get_user())
            elif command.command_type == CommandType.LOAD:
                response = self.api_client.load_bytes(command.filename)
                self.remote_handler.send(command.client_id, response)
            elif command.command_type == CommandType.INFERENCE:
                self.command_queue.put(command)
            elif command.command_type == CommandType.STOP_INFERENCE:
                self.inference.stop()
            elif command.command_type == CommandType.EXIT:
                self.stop()
            else:
                pass
        except Exception as e:
            print(f"Error handling client: {e}")

    cdef get_user(self, RemoteCommand command, User user):
        self.remote_handler.send(command.client_id, user.serialize())

    cdef on_annotation(self, RemoteCommand cmd, Annotation annotation):
        data = annotation.serialize()
        self.remote_handler.send(cmd.client_id, data)

    def stop(self):
        self.remote_handler.stop()
        self.running = False