77 lines
2.4 KiB
Python
Executable File
77 lines
2.4 KiB
Python
Executable File
import os, sys
|
|
cur_dir = os.path.dirname(__file__)
|
|
KIE_PATH = os.path.join(os.path.dirname(cur_dir), "sdsvkie")
|
|
TD_PATH = os.path.join(os.path.dirname(cur_dir), "sdsvtd")
|
|
TR_PATH = os.path.join(os.path.dirname(cur_dir), "sdsvtr")
|
|
sys.path.append(KIE_PATH)
|
|
sys.path.append(TD_PATH)
|
|
sys.path.append(TR_PATH)
|
|
|
|
from sdsvkie import Predictor
|
|
from .AnyKey_Value.anyKeyValue import load_engine, Predictor_KVU
|
|
import cv2
|
|
import numpy as np
|
|
import urllib
|
|
|
|
model = Predictor(
|
|
cfg = "/models/Kie_invoice_ap/06062023/config.yaml", # TODO: Better be scalable
|
|
device = "cuda:0",
|
|
weights = "/models/Kie_invoice_ap/06062023/best" # TODO: Better be scalable
|
|
)
|
|
|
|
class_names = ['others', 'title', 'key', 'value', 'header']
|
|
save_dir = os.path.join(cur_dir, "AnyKey_Value/visualize/test")
|
|
|
|
predictor, processor = load_engine(exp_dir="/models/Kie_invoice_fi/key_value_understanding-20230627-164536",
|
|
class_names=class_names,
|
|
mode=3)
|
|
|
|
def predict_fi(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"]
|
|
output_kie = {
|
|
field_name: field_item['value'] for field_name, field_item in output.items()
|
|
}
|
|
# print("Hoangggggggggggggggggggggggggggggggggggggggggggggg")
|
|
# print(output_kie)
|
|
|
|
|
|
#Phan cua Tuan
|
|
kvu_result, _ = Predictor_KVU(image_url, save_dir, predictor, processor)
|
|
# print("TuanNnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnn")
|
|
# print(kvu_result)
|
|
# if kvu_result['imei_number'] == None and kvu_result['serial_number'] == None:
|
|
return kvu_result, output_kie
|
|
|
|
if __name__ == "__main__":
|
|
image_url = "/mnt/hdd2T/dxtan/TannedCung/OCR/workspace/Kie_Invoice_AP/tmp_image/{image_url}.jpg"
|
|
output = predict_fi(0, image_url)
|
|
print(output) |