sbt-idp/cope2n-ai-fi/celery_worker/client_connector.py
2023-12-13 18:43:10 +07:00

82 lines
2.7 KiB
Python
Executable File

from celery import Celery
import base64
import environ
env = environ.Env(
DEBUG=(bool, True)
)
class CeleryConnector:
task_routes = {
"process_id_result": {"queue": "id_card_rs"},
"process_driver_license_result": {"queue": "driver_license_rs"},
"process_invoice_result": {"queue": "invoice_rs"},
"process_ocr_with_box_result": {"queue": "ocr_with_box_rs"},
"process_template_matching_result": {"queue": "template_matching_rs"},
# mock task
"process_id": {"queue": "id_card"},
"process_driver_license": {"queue": "driver_license"},
"process_invoice": {"queue": "invoice"},
"process_ocr_with_box": {"queue": "ocr_with_box"},
"process_template_matching": {"queue": "template_matching"},
}
app = Celery(
"postman",
broker=env.str("CELERY_BROKER", "amqp://test:test@rabbitmq:5672"),
broker_transport_options={'confirm_publish': True},
)
def process_id_result(self, args):
return self.send_task("process_id_result", args)
def process_driver_license_result(self, args):
return self.send_task("process_driver_license_result", args)
def process_invoice_result(self, args):
return self.send_task("process_invoice_result", args)
def process_ocr_with_box_result(self, args):
return self.send_task("process_ocr_with_box_result", args)
def process_template_matching_result(self, args):
return self.send_task("process_template_matching_result", args)
def process_id(self, args):
return self.send_task("process_id", args)
def process_driver_license(self, args):
return self.send_task("process_driver_license", args)
def process_invoice(self, args):
return self.send_task("process_invoice", args)
def process_ocr_with_box(self, args):
return self.send_task("process_ocr_with_box", args)
def process_template_matching(self, args):
return self.send_task("process_template_matching", args)
def send_task(self, name=None, args=None):
if name not in self.task_routes or "queue" not in self.task_routes[name]:
return self.app.send_task(name, args)
return self.app.send_task(name, args, queue=self.task_routes[name]["queue"])
def main():
rq_id = 345
file_names = "abc.jpg"
list_data = []
with open("/home/sds/thucpd/aicr-2022/abc.jpg", "rb") as fs:
encoded_string = base64.b64encode(fs.read()).decode("utf-8")
list_data.append(encoded_string)
c_connector = CeleryConnector()
a = c_connector.process_id(args=(rq_id, list_data, file_names))
print(a)
if __name__ == "__main__":
main()