PP-OCRv6_medium_det/src/example.py

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']}")