mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 08:36:29 +00:00
Alex Bezdieniezhnykh
98 lines
3.6 KiB
Cython
98 lines
3.6 KiB
Cython
# cython: language_level=3
|
|
import queue
|
|
import threading
|
|
import constants
|
|
from api_client import ApiClient
|
|
from inference import Inference, Annotation
|
|
from processor_command import FileCommand, CommandType, ProcessorType
|
|
from remote_handlers import SocketHandler, RabbitHandler
|
|
import argparse
|
|
|
|
cdef enum ListenOption:
|
|
SOCKET = 1
|
|
QUEUE = 2
|
|
|
|
cdef class ParsedArguments:
|
|
cdef ListenOption listen
|
|
cdef str email
|
|
cdef str password
|
|
cdef str folder
|
|
cdef bint persist_token
|
|
|
|
def __init__(self, ListenOption listen, str email, str password, str folder, bint persist_token):
|
|
self.listen = listen
|
|
self.email = email
|
|
self.password = password
|
|
self.folder = folder
|
|
self.persist_token = persist_token
|
|
|
|
cdef class CommandProcessor:
|
|
|
|
def __init__(self, args: ParsedArguments):
|
|
self.api_client = ApiClient(args.email, args.password, args.folder)
|
|
self.socket_handler = SocketHandler(self.on_message)
|
|
self.rabbit_handler = RabbitHandler(self.on_message)
|
|
self.command_queue = queue.Queue(maxsize=constants.QUEUE_MAXSIZE)
|
|
self.running = True
|
|
|
|
def start(self):
|
|
threading.Thread(target=self.process_queue, daemon=True).start()
|
|
|
|
cdef on_message(self, cmd: FileCommand):
|
|
try:
|
|
if cmd.command_type == CommandType.INFERENCE:
|
|
self.command_queue.put(cmd)
|
|
elif cmd.command_type == CommandType.LOAD:
|
|
threading.Thread(target=self.process_load, args=[cmd], daemon=True).start()
|
|
except Exception as e:
|
|
print(f"Error handling client: {e}")
|
|
|
|
cdef on_annotations(self, cmd: FileCommand, annotations: [Annotation]):
|
|
handler = self.socket_handler if cmd.processor_type == ProcessorType.SOCKET else self.rabbit_handler
|
|
handler.send(annotations)
|
|
|
|
|
|
cdef process_queue(self):
|
|
while self.running:
|
|
try:
|
|
command = self.command_queue.get()
|
|
model = self.api_client.load_ai_model()
|
|
Inference(model, self.on_annotations).run_inference(command)
|
|
except Exception as e:
|
|
print(f"Error processing queue: {e}")
|
|
|
|
cdef process_load(self, command: FileCommand):
|
|
response = self.api_client.load_file(command)
|
|
handler = self.socket_handler if command.processor_type == ProcessorType.SOCKET else self.rabbit_handler
|
|
handler.send(response)
|
|
|
|
|
|
def stop(self):
|
|
self.running = False
|
|
|
|
def parse_arguments():
|
|
parser = argparse.ArgumentParser(description="Command Processor")
|
|
parser.add_argument("--listen", type=ListenOption, choices=[ListenOption.SOCKET, ListenOption.QUEUE], default=ListenOption.SOCKET, help="socket: Local communication, queue: remote. Default is socket")
|
|
parser.add_argument("--email", type=str, default="", help="Email")
|
|
parser.add_argument("--pw", type=str, default="", help="Password")
|
|
parser.add_argument("--folder", type=str, default="", help="Folder to API inner folder to download file from")
|
|
parser.add_argument("--persist_token", type=bool, default=True, help="True for persisting token from API")
|
|
cdef args = parser.parse_args()
|
|
|
|
cdef ListenOption listen = ListenOption(args.listen)
|
|
cdef str email = args.email
|
|
cdef str password = args.pw
|
|
cdef str folder = args.folder
|
|
cdef bint persist_token = args.persist_token
|
|
|
|
return ParsedArguments(listen, email, password, folder, persist_token)
|
|
|
|
if __name__ == '__main__':
|
|
args = parse_arguments()
|
|
processor = CommandProcessor(args)
|
|
try:
|
|
processor.start()
|
|
except KeyboardInterrupt:
|
|
processor.stop()
|
|
|