File size: 2,265 Bytes
e4c2198 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | 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']}")
|