sbt-idp/cope2n-ai-fi/celery_worker/mock_process_tasks.py

220 lines
8.5 KiB
Python
Raw Normal View History

2023-11-30 11:22:16 +00:00
from celery_worker.worker import app
import numpy as np
import cv2
@app.task(name="process_id")
def process_id(rq_id, sub_id, folder_name, list_url, user_id):
from common.serve_model import predict
from celery_worker.client_connector import CeleryConnector
c_connector = CeleryConnector()
try:
result = predict(rq_id, sub_id, folder_name, list_url, user_id, infer_name="id_card")
print(result)
result = {
"status": 200,
"content": result,
"message": "Success",
}
c_connector.process_id_result((rq_id, result))
return {"rq_id": rq_id}
# if image_croped is not None:
# if result["data"] == []:
# result = {
# "status": 404,
# "content": {},
# }
# c_connector.process_id_result((rq_id, result, None))
# return {"rq_id": rq_id}
# else:
# result = {
# "status": 200,
# "content": result,
# "message": "Success",
# }
# c_connector.process_id_result((rq_id, result))
# return {"rq_id": rq_id}
# elif image_croped is None:
# result = {
# "status": 404,
# "content": {},
# }
# c_connector.process_id_result((rq_id, result, None))
# return {"rq_id": rq_id}
except Exception as e:
print(e)
result = {
"status": 404,
"content": {},
}
c_connector.process_id_result((rq_id, result, None))
return {"rq_id": rq_id}
@app.task(name="process_driver_license")
def process_driver_license(rq_id, sub_id, folder_name, list_url, user_id):
from common.serve_model import predict
from celery_worker.client_connector import CeleryConnector
c_connector = CeleryConnector()
try:
result = predict(rq_id, sub_id, folder_name, list_url, user_id, infer_name="driving_license")
result = {
"status": 200,
"content": result,
"message": "Success",
}
c_connector.process_driver_license_result((rq_id, result))
return {"rq_id": rq_id}
# result, image_croped = predict(str(url), "driving_license")
# if image_croped is not None:
# if result["data"] == []:
# result = {
# "status": 404,
# "content": {},
# }
# c_connector.process_driver_license_result((rq_id, result, None))
# return {"rq_id": rq_id}
# else:
# result = {
# "status": 200,
# "content": result,
# "message": "Success",
# }
# path_image_croped = "/app/media/users/{}/subscriptions/{}/requests/{}/{}/image_croped.jpg".format(user_id,sub_id,folder_name,rq_id)
# cv2.imwrite("/users/{}/subscriptions/{}/requests/{}/{}/image_croped.jpg".format(user_id,sub_id,folder_name,rq_id), image_croped)
# c_connector.process_driver_license_result((rq_id, result, path_image_croped))
# return {"rq_id": rq_id}
# elif image_croped is None:
# result = {
# "status": 404,
# "content": {},
# }
# c_connector.process_driver_license_result((rq_id, result, None))
# return {"rq_id": rq_id}
except Exception as e:
print(e)
result = {
"status": 404,
"content": {},
}
c_connector.process_driver_license_result((rq_id, result, None))
return {"rq_id": rq_id}
@app.task(name="process_template_matching")
def process_template_matching(rq_id, sub_id, folder_name, url, tmp_json, user_id):
from TemplateMatching.src.ocr_master import Extractor
from celery_worker.client_connector import CeleryConnector
import urllib
c_connector = CeleryConnector()
extractor = Extractor()
try:
req = urllib.request.urlopen(url)
arr = np.asarray(bytearray(req.read()), dtype=np.uint8)
img = cv2.imdecode(arr, -1)
imgs = [img]
image_aliged = extractor.image_alige(imgs, tmp_json)
if image_aliged is None:
result = {
"status": 401,
"content": "Image is not match with template",
}
c_connector.process_template_matching_result(
(rq_id, result, None)
)
return {"rq_id": rq_id}
else:
output = extractor.extract_information(
image_aliged, tmp_json
)
path_image_croped = "/app/media/users/{}/subscriptions/{}/requests/{}/{}/image_croped.jpg".format(user_id,sub_id,folder_name,rq_id)
cv2.imwrite("/users/{}/subscriptions/{}/requests/{}/{}/image_croped.jpg".format(user_id,sub_id,folder_name,rq_id), image_aliged)
if output == {}:
result = {"status": 404, "content": {}}
c_connector.process_template_matching_result((rq_id, result, None))
return {"rq_id": rq_id}
else:
result = {
"document_type": "template_matching",
"fields": []
}
print(output)
for field in tmp_json["fields"]:
print(field["label"])
field_value = {
"label": field["label"],
"value": output[field["label"]],
"box": [float(num) for num in field["box"]],
"confidence": 0.98 #TODO confidence
}
result["fields"].append(field_value)
print(result)
result = {"status": 200, "content": result}
c_connector.process_template_matching_result(
(rq_id, result, path_image_croped)
)
return {"rq_id": rq_id}
except Exception as e:
print(e)
result = {"status": 404, "content": {}}
c_connector.process_template_matching_result((rq_id, result, None))
return {"rq_id": rq_id}
# @app.task(name="process_invoice")
# def process_invoice(rq_id, url):
# from celery_worker.client_connector import CeleryConnector
# from Kie_Hoanglv.prediction import predict
# c_connector = CeleryConnector()
# try:
# print(url)
# result = predict(str(url))
# hoadon = {"status": 200, "content": result, "message": "Success"}
# c_connector.process_invoice_result((rq_id, hoadon))
# return {"rq_id": rq_id}
# except Exception as e:
# print(e)
# hoadon = {"status": 404, "content": {}}
# c_connector.process_invoice_result((rq_id, hoadon))
# return {"rq_id": rq_id}
@app.task(name="process_invoice")
def process_invoice(rq_id, list_url):
from celery_worker.client_connector import CeleryConnector
from common.process_pdf import compile_output
c_connector = CeleryConnector()
try:
result = compile_output(list_url)
hoadon = {"status": 200, "content": result, "message": "Success"}
c_connector.process_invoice_result((rq_id, hoadon))
return {"rq_id": rq_id}
except Exception as e:
print(e)
hoadon = {"status": 404, "content": {}}
c_connector.process_invoice_result((rq_id, hoadon))
return {"rq_id": rq_id}
@app.task(name="process_ocr_with_box")
def process_ocr_with_box(rq_id, list_url):
from celery_worker.client_connector import CeleryConnector
from common.process_pdf import compile_output_ocr_base
c_connector = CeleryConnector()
try:
result = compile_output_ocr_base(list_url)
result = {"status": 200, "content": result, "message": "Success"}
c_connector.process_ocr_with_box_result((rq_id, result))
return {"rq_id": rq_id}
except Exception as e:
print(e)
result = {"status": 404, "content": {}}
c_connector.process_ocr_with_box_result((rq_id, result))
return {"rq_id": rq_id}