import sys
import cv2
import numpy

sys.dont_write_bytecode = True

# Source
from helper import onnxSessionBuild

pathModel = "./PP-LCNet_x1_0_table_cls/"

imageSizeShort = 256
imageSizeCrop = 224

meanList = numpy.array([0.485, 0.456, 0.406], dtype=numpy.float32)
standardList = numpy.array([0.229, 0.224, 0.225], dtype=numpy.float32)

labelList = ["wired", "wireless"]

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

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

    ratio = imageSizeShort / float(min(imageHeight, imageWidth))

    imageResized = cv2.resize(imageRgb, (int(round(imageWidth * ratio)), int(round(imageHeight * ratio))))

    resizedHeight, resizedWidth = imageResized.shape[0:2]

    cropX = int(round((resizedWidth - imageSizeCrop) / 2))
    cropY = int(round((resizedHeight - imageSizeCrop) / 2))

    return imageResized[cropY:cropY + imageSizeCrop, cropX:cropX + imageSizeCrop]

def inference(imageRgb):
    imageCrop = imageResize(imageRgb)

    tensor = imageCrop.astype(numpy.float32) / 255.0
    tensor = (tensor - meanList) / standardList
    tensor = numpy.expand_dims(tensor.transpose((2, 0, 1)), axis=0).astype(numpy.float32)

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

    probabilityList = tensorOutputList[0][0]

    index = int(numpy.argmax(probabilityList))

    return {
        "label": labelList[index],
        "score": float(probabilityList[index])
    }

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

itemObject = inference(imageRgb)

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