File size: 2,465 Bytes
6c4ad51 a2f8908 6c4ad51 | 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 | 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}")
|