import sys
sys.dont_write_bytecode = True

import cv2
import numpy

from helper import onnxSessionBuild

pathModel = "./PP-DocLayout_plus-L/"

scoreThreshold = 0.3

labelObject = {
    12: "header",
    10: "doc_title",
    4: "abstract",
    5: "content",
    0: "paragraph_title",
    2: "text",
    1: "image",
    6: "figure_title",
    16: "chart",
    8: "table",
    7: "formula",
    17: "formula_number",
    13: "algorithm",
    18: "aside_text",
    9: "reference",
    19: "reference_content",
    11: "footnote",
    14: "footer",
    3: "number",
    15: "seal"
}

onnxSession = onnxSessionBuild(f"{pathModel}onnx/pp-docLayout_plus-l.onnx")

def inference(imageRgb):
    resultList = []

    imageHeight, imageWidth = imageRgb.shape[0:2]
    imageResized = cv2.resize(imageRgb, (800, 800), interpolation=cv2.INTER_CUBIC).astype(numpy.float32) / 255.0

    tensor = numpy.expand_dims(imageResized.transpose((2, 0, 1)), axis=0).astype(numpy.float32)

    tensorFeedObject = {
        "image": tensor,
        "im_shape": numpy.array([[800, 800]], dtype=numpy.float32),
        "scale_factor": numpy.array([[800 / float(imageHeight), 800 / float(imageWidth)]], dtype=numpy.float32)
    }

    tensorOutputList = onnxSession.run(None, tensorFeedObject)

    boxCount = int(tensorOutputList[1][0]) if len(tensorOutputList) > 1 else len(tensorOutputList[0])

    for a in range(boxCount):
        value = tensorOutputList[0][a]

        classId = int(value[0])
        score = float(value[1])
        x1 = max(0.0, min(float(value[2]), float(imageWidth)))
        y1 = max(0.0, min(float(value[3]), float(imageHeight)))
        x2 = max(0.0, min(float(value[4]), float(imageWidth)))
        y2 = max(0.0, min(float(value[5]), float(imageHeight)))

        if score >= scoreThreshold and x2 > x1 and y2 > y1:
            label = labelObject[classId] if classId in labelObject else str(classId)

            resultList.append({
                "label": label,
                "score": score,
                "coordinate": [x1, y1, x2, y2]
            })

    return resultList

image = cv2.imread(sys.argv[1])
imageRgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

itemList = inference(imageRgb)

for a in range(len(itemList)):
    coordinateList = []

    for b in range(len(itemList[a]["coordinate"])):
        coordinateList.append(int(round(itemList[a]["coordinate"][b])))

    print(f"{itemList[a]['score']:.6f} | {itemList[a]['label']} | {coordinateList}")