您的当前位置:首页正文

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

2024-12-01 来源:个人技术集锦

一、说明

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

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

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

进去,然后看排行榜,有些参赛者会分享训练好的模型。
说明:这些模型可能在自己数据上表现不是很好,因为用的很多都是合成图片训练的。

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

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

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

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

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

2.1 标注

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

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

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

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

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

2.2 标注文件处理

那么从labelme生成的JSON标注文件怎么转成yolo支持的像上面的形式呢。
我是先转成coco形式再转为yolo形式的,因为之前也有按照了中的方法训练了一下,这里的数据格式一般是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()

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

三、训练YOLOv8-pose

yolov8的项目

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

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

from ultralytics import YOLO

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

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

四、使用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)

五、小结

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

显示全文