如题,我的工作很多与OCR有关,主要就是针对卡证文档等图片进行OCR推理,我发现很多时候需要识别的信息因为拍摄或者扫描问题,导致信息与图片不在同一个水平面上,然后就需要一个预处理——透视矫正。
这里我先推荐两个我用过的开源方法
进去,然后看排行榜,有些参赛者会分享训练好的模型。
说明:这些模型可能在自己数据上表现不是很好,因为用的很多都是合成图片训练的。
进去,然后按照介绍上的modelscope的本地推理方法使用就可以了。
说明:这个模型矫正效果明显好了不少,但是只是针对卡证边缘进行了矫正,有些信息如果在一页文档但是内容是弯曲的,还是矫正不到位,这样需要用到OpenCV的透视变换。
因为卡证文档类型样式多样,不可能一个模型就可以通用所有场景,所以要想达到自己的要求,还是在垂类场景下微调模型吧。
YOLOv8-Pose属于TOP-Down方法,这个应该也属于两阶段的,其实就是先检测主体,然后在主体中检测关键点。可以参考
因为YOLOv8-Pose原本专门用于人体姿态估计的,需要标注一个主体检测框和17个人体关键点,而我想用在卡证文档上所需信息集中的区域作为主体,然后可以标注区域的四个角点或者再加上四个角点之间的中点,即八个点,这样就减少了不少功夫。
下面我是介绍标注四个角点的,八个角点其实也是差不多的方法。
关于YOLO的关键点检测我建议是使用labelme来进行,因为除了要标注关键信息区域,还要进行4个关键点的标注,那么4个关键点的标注就不能单纯用矩形框来标注了,需要使用能标注点的标注工具。
这里我用一张网上找到户口本图片作为示例,需要标注的是:
标注完成后保存会生成一个与图片同名的JSON格式的标注文件,上面展示部分内容。
最后应该会得到多个标注文件,那么接下来需要把这些标注文件转换成yolo能够读取的文件。
那么从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的项目
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", ……)
得到的模型就能对卡证文档图片进行关键点检测了。
最后就是进行透视矫正了,使用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的前置工作之一的图片透视矫正就是这样完成,经过测试,这个效果还是可以的。