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)