Files
annotations/Azaion.Inference/main.pyx
T
Alex Bezdieniezhnykh 0237e279a5 don't send user dto back
2025-04-17 00:18:40 +03:00

79 lines
2.8 KiB
Cython

import queue
import traceback
from queue import Queue
cimport constants
from threading import Thread
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 credentials cimport Credentials
from file_data cimport FileData
from user cimport User
cdef class CommandProcessor:
cdef ApiClient api_client
cdef RemoteCommandHandler remote_handler
cdef object inference_queue
cdef bint running
cdef Inference inference
def __init__(self):
self.api_client = ApiClient()
self.remote_handler = RemoteCommandHandler(self.on_command)
self.inference_queue = Queue(maxsize=constants.QUEUE_MAXSIZE)
self.remote_handler.start()
self.running = True
self.inference = Inference(self.api_client, self.on_annotation)
def start(self):
while self.running:
try:
command = self.inference_queue.get(timeout=0.5)
self.inference.run_inference(command)
self.remote_handler.send(command.client_id, <bytes>'DONE'.encode('utf-8'))
except queue.Empty:
continue
except Exception as e:
traceback.print_exc()
print('EXIT!')
cdef on_command(self, RemoteCommand command):
try:
if command.command_type == CommandType.LOGIN:
self.login(command)
elif command.command_type == CommandType.LOAD:
self.load_file(command)
elif command.command_type == CommandType.INFERENCE:
self.inference_queue.put(command)
elif command.command_type == CommandType.STOP_INFERENCE:
self.inference.stop()
elif command.command_type == CommandType.EXIT:
t = Thread(target=self.stop) # non-block worker:
t.start()
else:
pass
except Exception as e:
print(f"Error handling client: {e}")
cdef login(self, RemoteCommand command):
cdef User user
self.api_client.set_credentials(Credentials.from_msgpack(command.data))
user = self.api_client.get_user()
cdef load_file(self, RemoteCommand command):
cdef FileData file_data = FileData.from_msgpack(command.data)
response = self.api_client.load_bytes(file_data.filename, file_data.folder)
self.remote_handler.send(command.client_id, response)
cdef on_annotation(self, RemoteCommand cmd, Annotation annotation):
data = annotation.serialize()
self.remote_handler.send(cmd.client_id, data)
def stop(self):
self.inference.stop()
self.remote_handler.stop()
self.running = False