mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 21:56:31 +00:00
dzaitsev
92 lines
3.7 KiB
Cython
92 lines
3.7 KiB
Cython
import queue
|
|
import traceback
|
|
from queue import Queue
|
|
cimport constants
|
|
from threading import Thread
|
|
|
|
import yaml
|
|
|
|
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
|
|
|
|
cdef class CommandProcessor:
|
|
cdef ApiClient api_client
|
|
cdef RemoteCommandHandler remote_handler
|
|
cdef object inference_queue
|
|
cdef bint running
|
|
cdef Inference inference
|
|
|
|
def __init__(self, int zmq_port, str api_url):
|
|
self.api_client = ApiClient(api_url)
|
|
self.remote_handler = RemoteCommandHandler(zmq_port, 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)
|
|
end_inference_command = RemoteCommand(CommandType.INFERENCE_DATA, None, 'DONE')
|
|
self.remote_handler.send(command.client_id, end_inference_command.serialize())
|
|
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.api_client.set_credentials(Credentials.from_msgpack(command.data))
|
|
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.AI_AVAILABILITY_CHECK:
|
|
self.build_tensor_engine(command.client_id)
|
|
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 build_tensor_engine(self, client_id):
|
|
self.inference.build_tensor_engine(lambda status: self.build_tensor_status_updater(client_id, status))
|
|
self.remote_handler.send(client_id, RemoteCommand(CommandType.AI_AVAILABILITY_RESULT, None, 'enabled').serialize())
|
|
|
|
cdef build_tensor_status_updater(self, bytes client_id, str status):
|
|
self.remote_handler.send(client_id, RemoteCommand(CommandType.AI_AVAILABILITY_RESULT, None, status).serialize())
|
|
|
|
cdef load_file(self, RemoteCommand command):
|
|
cdef RemoteCommand response
|
|
cdef FileData file_data
|
|
cdef bytes file_bytes
|
|
try:
|
|
file_data = FileData.from_msgpack(command.data)
|
|
file_bytes = self.api_client.load_bytes(file_data.filename, file_data.folder)
|
|
response = RemoteCommand(CommandType.DATA_BYTES, file_bytes)
|
|
except Exception as e:
|
|
response = RemoteCommand(CommandType.DATA_BYTES, None, str(e))
|
|
self.remote_handler.send(command.client_id, response.serialize())
|
|
|
|
cdef on_annotation(self, RemoteCommand cmd, Annotation annotation):
|
|
cdef RemoteCommand response = RemoteCommand(CommandType.INFERENCE_DATA, annotation.serialize())
|
|
self.remote_handler.send(cmd.client_id, response.serialize())
|
|
|
|
def stop(self):
|
|
self.inference.stop()
|
|
self.remote_handler.stop()
|
|
self.running = False
|