import sys
sys.dont_write_bytecode = True

import io
import math
import cv2
import numpy

from helper import onnxSessionBuild

pathModel = "./PP-OCRv6_medium_rec/"

imageHeightModel = 48
imageWidthModel = 320
imageWidthMax = 3200

characterList = ["blank"]

lineList = io.open(f"{pathModel}onnx/dictionary.txt", encoding="utf-8").read().split("\n")

for a in range(len(lineList)):
    if lineList[a] != "":
        characterList.append(lineList[a])

characterList.append(" ")

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

def imageResize(image):
    imageHeight, imageWidth = image.shape[0:2]

    ratioWidthHeight = max(imageWidthModel / float(imageHeightModel), imageWidth / float(imageHeight))

    widthTarget = int(imageHeightModel * ratioWidthHeight)

    if widthTarget > imageWidthMax:
        widthTarget = imageWidthMax
        widthResized = imageWidthMax
    else:
        widthResized = int(math.ceil(imageHeightModel * imageWidth / float(imageHeight)))

        if widthResized > widthTarget:
            widthResized = widthTarget

    imageResized = cv2.resize(image, (widthResized, imageHeightModel))

    tensor = imageResized.astype(numpy.float32).transpose((2, 0, 1)) / 255.0
    tensor = (tensor - 0.5) / 0.5

    tensorPadded = numpy.zeros((3, imageHeightModel, widthTarget), dtype=numpy.float32)
    tensorPadded[:, :, 0:widthResized] = tensor

    return numpy.expand_dims(tensorPadded, axis=0)

def inference(image):
    tensor = imageResize(image)

    tensorOutputList = onnxSession.run(None, {"x": tensor})

    probability = tensorOutputList[0][0]

    indexList = probability.argmax(axis=-1)
    valueList = probability.max(axis=-1)

    text = ""
    scoreList = []

    for a in range(len(indexList)):
        if indexList[a] == 0:
            continue

        if a > 0 and indexList[a] == indexList[a - 1]:
            continue

        text += characterList[indexList[a]]

        scoreList.append(float(valueList[a]))

    score = 0.0

    if len(scoreList) > 0:
        score = float(numpy.mean(scoreList))

    return {
        "text": text,
        "score": score
    }

image = cv2.imread(sys.argv[1])

itemObject = inference(image)

print(f"{itemObject['score']:.6f} | {itemObject['text']}")