56 lines
1.5 KiB
Python
Executable File
56 lines
1.5 KiB
Python
Executable File
from sdsvkie import Predictor
|
|
import cv2
|
|
import numpy as np
|
|
import urllib
|
|
|
|
model = Predictor(
|
|
cfg = "/ai-core/models/Kie_invoice_ap/config.yaml",
|
|
device = "cuda:0",
|
|
weights = "/ai-core/models/Kie_invoice_ap/ep21"
|
|
)
|
|
|
|
def predict(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)
|
|
out = model(img)
|
|
output = out["end2end_results"]
|
|
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']
|
|
}
|
|
output_dict['fields'].append(field)
|
|
return output_dict
|
|
|
|
if __name__ == "__main__":
|
|
image_url = "/mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/demos/2022_07_25 farewell lunch.jpg"
|
|
output = predict(image_url) |