update

PP-OCRv6_medium_det/src/example.py

@@ -0,0 +1,162 @@
+import sys
+sys.dont_write_bytecode = True
+
+import cv2
+import numpy
+import pyclipper
+
+from helper import onnxSessionBuild
+
+pathModel = "./PP-OCRv6_medium_det/"
+
+limitSideLength = 960
+limitSideLengthMax = 4000
+binaryThreshold = 0.2
+boxThreshold = 0.45
+unclipRatio = 1.4
+candidateMax = 3000
+sideMinimum = 3
+
+meanList = numpy.array([0.485, 0.456, 0.406], dtype=numpy.float32)
+standardList = numpy.array([0.229, 0.224, 0.225], dtype=numpy.float32)
+
+onnxSession = onnxSessionBuild(f"{pathModel}onnx/pp-ocrV6_medium_det.onnx")
+
+def imageResize(image):
+    imageHeight, imageWidth = image.shape[0:2]
+
+    ratio = 1.0
+
+    if max(imageHeight, imageWidth) > limitSideLength:
+        if imageHeight > imageWidth:
+            ratio = float(limitSideLength) / imageHeight
+        else:
+            ratio = float(limitSideLength) / imageWidth
+
+    resizeHeight = int(imageHeight * ratio)
+    resizeWidth = int(imageWidth * ratio)
+
+    if max(resizeHeight, resizeWidth) > limitSideLengthMax:
+        ratio = float(limitSideLengthMax) / max(resizeHeight, resizeWidth)
+
+        resizeHeight = int(resizeHeight * ratio)
+        resizeWidth = int(resizeWidth * ratio)
+
+    resizeHeight = max(int(round(resizeHeight / 32) * 32), 32)
+    resizeWidth = max(int(round(resizeWidth / 32) * 32), 32)
+
+    return cv2.resize(image, (resizeWidth, resizeHeight))
+
+def boxOrder(contour):
+    rectangle = cv2.minAreaRect(contour)
+
+    pointList = sorted(list(cv2.boxPoints(rectangle)), key=lambda point: point[0])
+
+    index1 = 0
+    index2 = 1
+    index3 = 2
+    index4 = 3
+
+    if pointList[1][1] > pointList[0][1]:
+        index1 = 0
+        index4 = 1
+    else:
+        index1 = 1
+        index4 = 0
+
+    if pointList[3][1] > pointList[2][1]:
+        index2 = 2
+        index3 = 3
+    else:
+        index2 = 3
+        index3 = 2
+
+    return [pointList[index1], pointList[index2], pointList[index3], pointList[index4]], min(rectangle[1])
+
+def boxScore(probabilityMap, box):
+    mapHeight, mapWidth = probabilityMap.shape[0:2]
+
+    boxLocal = box.copy()
+
+    xMinimum = max(0, min(int(numpy.floor(box[:, 0].min())), mapWidth - 1))
+    xMaximum = max(0, min(int(numpy.ceil(box[:, 0].max())), mapWidth - 1))
+    yMinimum = max(0, min(int(numpy.floor(box[:, 1].min())), mapHeight - 1))
+    yMaximum = max(0, min(int(numpy.ceil(box[:, 1].max())), mapHeight - 1))
+
+    mask = numpy.zeros((yMaximum - yMinimum + 1, xMaximum - xMinimum + 1), dtype=numpy.uint8)
+
+    boxLocal[:, 0] = boxLocal[:, 0] - xMinimum
+    boxLocal[:, 1] = boxLocal[:, 1] - yMinimum
+
+    cv2.fillPoly(mask, boxLocal.reshape(1, -1, 2).astype(numpy.int32), 1)
+
+    return cv2.mean(probabilityMap[yMinimum:yMaximum + 1, xMinimum:xMaximum + 1], mask)[0]
+
+def boxUnclip(box):
+    area = cv2.contourArea(box)
+    length = cv2.arcLength(box, True)
+
+    distance = area * unclipRatio / length
+
+    offsetObject = pyclipper.PyclipperOffset()
+
+    offsetObject.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+
+    return numpy.array(offsetObject.Execute(distance))
+
+def inference(image):
+    resultList = []
+
+    imageHeight, imageWidth = image.shape[0:2]
+    imageResized = imageResize(image)
+
+    tensor = (imageResized.astype(numpy.float32) / 255.0 - meanList) / standardList
+    tensor = numpy.expand_dims(tensor.transpose((2, 0, 1)), axis=0).astype(numpy.float32)
+
+    tensorOutputList = onnxSession.run(None, {"x": tensor})
+
+    probabilityMap = tensorOutputList[0][0][0]
+    bitmap = (probabilityMap > binaryThreshold).astype(numpy.uint8)
+
+    scaleWidth = imageWidth / float(bitmap.shape[1])
+    scaleHeight = imageHeight / float(bitmap.shape[0])
+
+    contourList, hierarchy = cv2.findContours(bitmap * 255, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+
+    for a in range(min(len(contourList), candidateMax)):
+        pointList, sideLength = boxOrder(contourList[a])
+
+        if sideLength < sideMinimum:
+            continue
+
+        score = boxScore(probabilityMap, numpy.array(pointList).reshape(-1, 2))
+
+        if score < boxThreshold:
+            continue
+
+        pointList, sideLength = boxOrder(boxUnclip(numpy.array(pointList)).reshape(-1, 1, 2))
+
+        if sideLength < sideMinimum + 2:
+            continue
+
+        coordinateList = []
+
+        for b in range(len(pointList)):
+            x = max(0, min(int(round(pointList[b][0] * scaleWidth)), imageWidth))
+            y = max(0, min(int(round(pointList[b][1] * scaleHeight)), imageHeight))
+
+            coordinateList.append([x, y])
+
+        resultList.append({
+            "score": score,
+            "coordinate": coordinateList
+        })
+
+    return resultList
+
+image = cv2.imread(sys.argv[1])
+
+itemList = inference(image)
+
+for a in range(len(itemList)):
+    print(f"{itemList[a]['score']:.6f} | {itemList[a]['coordinate']}")