TPS薄板样条变换

之前使用TPS（Thin-Plate Spline）对弯曲的文本进行处理，进而识别处理之后的水平文本，记录一下备忘。

话不多说，直接使用CV2 的 TPS Transformer：

import json
import cv2
import math
import numpy as np

with open("Label.txt","r",encoding='utf-8') as f:
    content = f.read()
contents = content.split("\n")
with open("testTps.txt","a",encoding='utf-8') as wf:
    for c in contents:
        if "\t" in c:
            tmp = c.split("\t")
            img = "./"+tmp[0]
            boxes = json.loads(tmp[1])
            imgori = cv2.imread(img)
            box_num = 0
            for box in boxes:
                #开始执行TPS变换，points是弯曲文本的标注，从左上角开始顺时针
                points = box["points"]
                #获取能把弯曲文本框包围住的矩形框
                x, y, w, h = cv2.boundingRect(np.array(points))
                # print([x,y,w,h])
                #分割出矩形框
                roi_img = imgori[y:y + h, x:x + w]
                # cv2.imwrite("./roiimg"+str(box_num)+".png",roi_img)
                #初始化Transformer
                tps = cv2.createThinPlateSplineShapeTransformer()

                sourceshape = np.array(points, np.int32)
                sourceshape = sourceshape.reshape(1, -1, 2)
                # print(sourceshape)
                matches = []
                N = len(points)
                for i in range(0, N):
                    matches.append(cv2.DMatch(i, i, 0))
                #最终结果的高度采用第一个点与最后一个点之间的距离
                h = int(math.sqrt(((points[0][0] - points[N - 1][0]) ** 2) +
                                    ((points[0][1] - points[N - 1][1]) ** 2)))
                # print("height is "+str(h))
                # 开始变动，计算变换后的目标点
                newpoints = []
                N = N // 2
                dx = int(w / (N - 1))
                for i in range(0, N):
                    newpoints.append((dx * i, 2))
                for i in range(N - 1, -1, -1):
                    newpoints.append((dx * i, h))
                # print(newpoints)
                targetshape = np.array(newpoints, np.int32)
                targetshape = targetshape.reshape(1, -1, 2)
                tps.estimateTransformation(targetshape, sourceshape, matches)
                # print(targetshape)
                try:
                    roi_img_ = tps.warpImage(imgori)
                except:
                    print("TPS变换错误")
                    break
                # cv2.imwrite("./roi_img_"+str(box_num)+".png",roi_img_)
                try:
                    img = roi_img_[newpoints[0][1]:newpoints[N][1], newpoints[0][0]:newpoints[N][0]]
                except:
                    print("TPS变换后无法存储")
                    break
                filename = "./tps/"+tmp[0].split(".")[0]+"-"+str(box_num)+".png"
                # print(filename)
                wf.write(tmp[0].split(".")[0]+"-"+str(box_num)+".png"+"\t"+box["transcription"]+"\n")
                cv2.imwrite(filename,img)
                box_num += 1

Label.txt的格式为PaddleOCR的标注格式：

seal/1-1.png	[{"transcription": "全国统一发票监制章", "points": [[23, 66], [35, 49], [54, 36], [80, 25], [107, 21], [137, 21], [162, 25], [184, 33], [207, 49], [220, 66], [203, 78], [189, 64], [175, 53], [155, 46], [135, 42], [109, 40], [84, 42], [64, 51], [50, 64], [40, 77]]}, {"transcription": "国家税务总局", "points": [[64, 73], [76, 73], [89, 73], [102, 73], [115, 73], [128, 73], [141, 73], [154, 73], [167, 73], [180, 73], [180, 94], [167, 94], [154, 94], [141, 94], [128, 94], [115, 94], [102, 94], [89, 94], [76, 94], [64, 94]]}, {"transcription": "山东省税务局", "points": [[71, 112], [81, 115], [91, 116], [102, 118], [115, 119], [129, 119], [139, 118], [149, 116], [157, 114], [172, 110], [178, 130], [164, 134], [154, 138], [143, 141], [131, 142], [115, 142], [102, 140], [87, 137], [76, 134], [65, 130]]}]

最终输出：

testTps.txt ：

seal/1-1-0.png	全国统一发票监制章
seal/1-1-1.png	国家税务总局
seal/1-1-2.png	山东省税务局