【OCR】OCR前置工作——透视矫正，使用YOLOv8-pose分支在卡证文档图片上进行yolo关键点检测_透视变换应用到yolo检测中

一、说明

如题，我的工作很多与OCR有关，主要就是针对卡证文档等图片进行OCR推理，我发现很多时候需要识别的信息因为拍摄或者扫描问题，导致信息与图片不在同一个水平面上，然后就需要一个预处理——透视矫正。

这里我先推荐两个我用过的开源方法

1.1 百度网盘AI大赛——表格检测参赛者分享的模型

进去《百度网盘AI大赛——表格检测》，然后看排行榜，有些参赛者会分享训练好的模型。
说明：这些模型可能在自己数据上表现不是很好，因为用的很多都是合成图片训练的。

2.2 阿里达摩院读光-票证检测矫正模型

进去《读光-票证检测矫正模型》，然后按照介绍上的modelscope的本地推理方法使用就可以了。
说明：这个模型矫正效果明显好了不少，但是只是针对卡证边缘进行了矫正，有些信息如果在一页文档但是内容是弯曲的，还是矫正不到位，这样需要用到OpenCV的透视变换。

二、针对自己数据训练一个YOLO-Pose模型

因为卡证文档类型样式多样，不可能一个模型就可以通用所有场景，所以要想达到自己的要求，还是在垂类场景下微调模型吧。
YOLOv8-Pose属于TOP-Down方法，这个应该也属于两阶段的，其实就是先检测主体，然后在主体中检测关键点。可以参考《yolov8-pose关键点检测，从数据集制作到训练测试》
因为YOLOv8-Pose原本专门用于人体姿态估计的，需要标注一个主体检测框和17个人体关键点，而我想用在卡证文档上所需信息集中的区域作为主体，然后可以标注区域的四个角点或者再加上四个角点之间的中点，即八个点，这样就减少了不少功夫。

下面我是介绍标注四个角点的，八个角点其实也是差不多的方法。

2.1 标注

关于YOLO的关键点检测我建议是使用labelme来进行，因为除了要标注关键信息区域，还要进行4个关键点的标注，那么4个关键点的标注就不能单纯用矩形框来标注了，需要使用能标注点的标注工具。

这里我用一张网上找到户口本图片作为示例，需要标注的是：

• box：关键信息区域主题，使用矩形框包围，要全部包围。
• tl：左上角点（图上面是lr，输入错了）
• tr：右上角点
• br：右下角点
• bl：左下角点
  

标注完成后保存会生成一个与图片同名的JSON格式的标注文件，上面展示部分内容。

最后应该会得到多个标注文件，那么接下来需要把这些标注文件转换成yolo能够读取的文件。

一共17个数据，第一个数据是0，表示该目标的类别，即box；后面跟着的四个数据，表示矩形框的坐标；接下来的是4*3， 表示4个关键点的坐标及是否可见。
其中 0.00000 表示没有显露出不可见，1.00000 表示被遮挡不可见，2.00000 表示可见。这个按图片情况具体修改。

2.2 标注文件处理

那么从labelme生成的JSON标注文件怎么转成yolo支持的像上面的形式呢。
我是先转成coco形式再转为yolo形式的，因为之前也有按照了《百度网盘AI大赛——表格检测》中的方法训练了一下，这里的数据格式一般是coco形式的。
其中有些小细节就不具体说了，网上还是有很多labelme导出标注文件转coco形式的，上面发的yolov8-pose流程参考中也有这个。下面给个使用coco转yolo形式的。

# 生成点坐标
def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = box[0] + box[2] / 2.0
    y = box[1] + box[3] / 2.0
    w = box[2]
    h = box[3]

    x = round(x * dw, 6)
    w = round(w * dw, 6)
    y = round(y * dh, 6)
    h = round(h * dh, 6)
    return (x, y, w, h)
    
# json_file：COCO Object Instance 类型的标注
# ana_txt_save_path：设置.txt文件保存位置

data = json.load(open(json_file, 'r'))
if not os.path.exists(ana_txt_save_path):
    os.makedirs(ana_txt_save_path)

id_map = {}  # coco数据集的id不连续，需要重新映射一下
with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
    # 写入classes.txt，这个其实只有一个“box”标签
    for i, category in enumerate(data['categories']):
        f.write(category['name']+"\n")
        id_map[category['id']] = i
    
for img in tqdm(data['images']):
        filename = img["file_name"]   # 获取图片文件名
        img_width = img["width"]		  # 获取图片的宽
        img_height = img["height"]	  # 获取图片的高
        img_id = img["id"] 				  # 获取图片id
        
        # yolo标注文件对应图片的txt名字，与jpg一致
        head, tail = os.path.splitext(filename) 
        ana_txt_name = head + ".txt"  	
        
        f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')
        for ann in data['annotations']:
            if ann['image_id'] == img_id:
                box = convert((img_width, img_height), ann["box"])
                f_txt.write("%s %s %s %s %s" % (id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))
                counter=0
                for i in range(len(ann["points"])):
                    if ann["points"][i] == 2 or ann["points"][i] == 1 or ann["points"][i] == 0:
                        f_txt.write(" %s " % format(ann["points"][i] + 1,'6f'))
                        counter=0
                    else:
                        if counter==0:
                            f_txt.write(" %s " % round((ann["points"][i] / img_width),6))
                        else:
                            f_txt.write(" %s " % round((ann["points"][i] / img_height),6))
                        counter+=1
        f_txt.write("\n")
        f_txt.close()
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

这样就生成yolo支持的关键点检测标注数据了。

三、训练YOLOv8-pose

yolov8的项目《ultralytics项目》

编辑提供YOLO读取的数据文件posedata.yaml，可以不设置标签文件路径，可以像我这样放，目录分别放图片和标签文件

下一级目录如下，其中标注目录记得放一个类别文件classes.txt，其中只有一类，即“box”

然后yolo读取的数据文件posedata.yaml内容如下：

train: ……/data/keypoint/images/train
val: ……/data/keypoint/images/val

# Keypoints
kpt_shape: [4, 3]  # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
flip_idx: [1,0,3,2] # 图像翻转后原关键点的ID对应原标签id的位置，比如0和1对称，原位置为[0,1]那么翻转后[1,0]

names:
  0: box
1
2
3
4
5
6
7
8
9

训练的话其实就是预训练模型换成pose分支的，比如

from ultralytics import YOLO

model = YOLO('yolov8m-pose.yaml').load('yolov8m-pose.pt')
model.train(data="posedata.yaml", ……)
1
2
3
4

得到的模型就能对卡证文档图片进行关键点检测了。

四、使用yolo关键点检测模型实现图片的矫正

最后就是进行透视矫正了，使用OpenCV实现，把所需信息区域矫正成在同一水平面。

point_result = keypoint.inference(img)
if point_result.shape[1] == 17: # 为什么是17可以看看标注文件txt
    for i in point_result:
        bbox, lt, rt, rb, lb = list(i[0:5]), list(i[5:8]), list(i[8:11]), list(i[11:14]), list(i[14:])
    if bbox[-1] > 0.7:
        point = [lt[:-1], rt[:-1], rb[:-1], lb[:-1]]
            
p1 = np.array(point, dtype=np.float32)	# point是yolo-pose检测出来的4个关键点(x,y)坐标
p2 = np.array([左上坐标(x,y)，右上坐标(x,y)，右下坐标(x,y)，左下坐标(x,y)], dtype=np.float32) # 这个看图片了，把检测的4个关键点拉到对应坐标
M = cv2.getPerspectiveTransform(p1,p2)
img = cv2.warpPerspective(img, M, (宽, 高), borderValue=(255, 255, 255))  # 将矩形图片映射为任意四边形
cv2.imwrite('test.jpg', img)
1
2
3
4
5
6
7
8
9
10
11
12

五、小结

对于ocr的前置工作之一的图片透视矫正就是这样完成，经过测试，这个效果还是可以的。