144 lines
4.6 KiB
Python
Executable File
144 lines
4.6 KiB
Python
Executable File
from sdsvkie import Predictor
|
|
import sys, os
|
|
cur_dir = os.path.dirname(__file__)
|
|
# sys.path.append("cope2n-ai-fi/Kie_Invoice_AP") # Better be relative
|
|
from .AnyKey_Value.anyKeyValue import load_engine, Predictor_KVU
|
|
import cv2
|
|
import numpy as np
|
|
import urllib
|
|
import random
|
|
|
|
# model = Predictor(
|
|
# cfg = "/cope2n-ai-fi/models/Kie_invoice_ap/config.yaml",
|
|
# device = "cuda:0",
|
|
# weights = "/cope2n-ai-fi/models/Kie_invoice_ap/best"
|
|
# )
|
|
|
|
class_names = ['others', 'title', 'key', 'value', 'header']
|
|
save_dir = os.path.join(cur_dir, "AnyKey_Value/visualize/test")
|
|
|
|
predictor, processor = load_engine(exp_dir=os.path.join(cur_dir, "AnyKey_Value/experiments/key_value_understanding-20231003-171748"),
|
|
class_names=class_names,
|
|
mode=3)
|
|
|
|
def predict(page_numb, image_url):
|
|
"""
|
|
module predict function
|
|
|
|
Args:
|
|
image_url (str): image url
|
|
|
|
Returns:
|
|
example output:
|
|
"data": {
|
|
"document_type": "invoice",
|
|
"fields": [
|
|
{
|
|
"label": "Invoice Number",
|
|
"value": "INV-12345",
|
|
"box": [0, 0, 0, 0],
|
|
"confidence": 0.98
|
|
},
|
|
...
|
|
]
|
|
}
|
|
dict: output of model
|
|
"""
|
|
req = urllib.request.urlopen(image_url)
|
|
arr = np.asarray(bytearray(req.read()), dtype=np.uint8)
|
|
img = cv2.imdecode(arr, -1)
|
|
# img = cv2.imread(image_url)
|
|
|
|
# Phan cua LeHoang
|
|
# out = model(img)
|
|
# output = out["end2end_results"]
|
|
|
|
|
|
#Phan cua Tuan
|
|
kvu_result = Predictor_KVU(image_url, save_dir, predictor, processor)
|
|
output_dict = {
|
|
"document_type": "invoice",
|
|
"fields": []
|
|
}
|
|
for key in kvu_result.keys():
|
|
field = {
|
|
"label": key,
|
|
"value": kvu_result[key],
|
|
"box": [0, 0, 0, 0],
|
|
"confidence": random.uniform(0.9, 1.0),
|
|
"page": page_numb
|
|
}
|
|
output_dict['fields'].append(field)
|
|
return output_dict
|
|
|
|
# if kvu_result['imei_number'] == None and kvu_result['serial_number'] == None:
|
|
# output_dict = {
|
|
# "document_type": "invoice",
|
|
# "fields": []
|
|
# }
|
|
# for key in output.keys():
|
|
# field = {
|
|
# "label": key if key != "id" else "Receipt Number",
|
|
# "value": output[key]['value'] if output[key]['value'] else "",
|
|
# "box": output[key]['box'],
|
|
# "confidence": output[key]['conf'],
|
|
# "page": page_numb
|
|
# }
|
|
# output_dict['fields'].append(field)
|
|
# table = kvu_result['table']
|
|
# field_table = {
|
|
# "label": "table",
|
|
# "value": table,
|
|
# "box": [0, 0, 0, 0],
|
|
# "confidence": 0.98,
|
|
# "page": page_numb
|
|
# }
|
|
# output_dict['fields'].append(field_table)
|
|
# return output_dict
|
|
|
|
# else:
|
|
# output_dict = {
|
|
# "document_type": "KSU",
|
|
# "fields": []
|
|
# }
|
|
# # for key in output.keys():
|
|
# # field = {
|
|
# # "label": key if key != "id" else "Receipt Number",
|
|
# # "value": output[key]['value'] if output[key]['value'] else "",
|
|
# # "box": output[key]['box'],
|
|
# # "confidence": output[key]['conf'],
|
|
# # "page": page_numb
|
|
# # }
|
|
# # output_dict['fields'].append(field)
|
|
|
|
# # Serial Number
|
|
# serial_number = kvu_result['serial_number']
|
|
# field_serial = {
|
|
# "label" : "serial_number",
|
|
# "value": serial_number,
|
|
# "box": [0, 0, 0, 0],
|
|
# "confidence": 0.98,
|
|
# "page": page_numb
|
|
# }
|
|
# output_dict['fields'].append(field_serial)
|
|
|
|
# # IMEI Number
|
|
# imei_number = kvu_result['imei_number']
|
|
# if imei_number == None:
|
|
# return output_dict
|
|
# if imei_number != None:
|
|
# for i in range(len(imei_number)):
|
|
# field_imei = {
|
|
# "label": "imei_number_{}".format(i+1),
|
|
# "value": imei_number[i],
|
|
# "box": [0, 0, 0, 0],
|
|
# "confidence": 0.98,
|
|
# "page": page_numb
|
|
# }
|
|
# output_dict['fields'].append(field_imei)
|
|
|
|
# return output_dict
|
|
|
|
if __name__ == "__main__":
|
|
image_url = "/root/thucpd/20230322144639VUzu_16794962527791962785161104697882.jpg"
|
|
output = predict(0, image_url) |