update

PP-OCRv6_medium_rec/src/example.py

@@ -0,0 +1,93 @@
+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']}")