【摘要】 这是水表读数识别项目,实现了如何端到端完成水表读数识别项目。涉及领域包括图像分类、语义分割、OCR文本检测、OCR文本识别。
水表读数识别项目简介:
这里实现了如何端到端完成水表读数识别项目。涉及领域包括图像分类、语义分割、OCR文本检测、OCR文本识别。
本案例提供的方法较多,涉及多个模型,但不需要运行所有的模型。如果都测试一遍,默认创建的5GB磁盘规格可能不够用。磁盘空间不够时可以将不需要的数据和模型文件删除,腾出空间;或者创建NoteBook时将磁盘规格增加到10GB。
解决方案流程:
一、使用语义分割或者OCR文本检测算法识别水表读数所在的四边形区域,并对四边形进行仿射变换转换成矩形,保存成新的图片数据。
二、如果文本区域的图片翻转严重,那第一步抠图生成的数据可能会有180度的翻转。因为文本识别的算法对翻转180度的场景效果不佳,所以再训练一个识别文本翻转的分类模型(本案例图片旋转角度微小,不需要此步骤,内容可供参考)。
三、利用步骤一中检测并抠图出来的文本数据训练OCR文本识别算法,识别图片中的文本内容,即数字。
数据集为华为云AI Gallery上提供的数据集。详情可参见本案例的关联资产。
下载该项目依赖脚本:
!wget --no-check-certificate https://modelarts-cnnorth4-market.obs.cn-north-4.myhuaweicloud.com/moxing-apps/notebooks/WaterMeter_Identification/deps.zip -O deps.zip!unzip -qo deps.zip!rm -f deps.zip!ls
准备环境
准备数据
数据下载
数据切分
文本检测
数据准备
下载预训练模型
训练
推理
抠图
数据准备
下载预训练模型
训练
推理 四边形定位 矫正
分割算法(推荐)
文本检测算法
文本翻转检测(可选)
数据准备
下载预训练模型
训练
推理
文本识别
数据准备
下载预训练模型
训练
推理
端到端识别
分割算法
文本识别算法
如果没有V100,也可以选择P100,如图所示,看看配置:
为了方便训练模型时观察模型的训练效果,将数据集切分成训练集和验证集,默认切分比例 训练集:验证集=9:1,保存在./data/train.txt和./data/eval.txt。
from util import split_dataimg_dir = './data/raw/train_imgs' # 原始图片路径save_dir = './data/' # 切分结果保存路径split_data(img_dir, save_dir) # 开始切分print('Finished!')
!ls ./data
本案例尝试了两种算法识别度数区域。分别是语义分割算法deeplabv3和文本检测算法dbnet。对比之后分割算法结果更佳。
分割算法和文本检测算法两者选择其一即可,推荐分割算法。
与文本检测算法二选一即可。
分割算法的数据标签通常为单通道的PNG图片。所以我们需要通过原始的txt标签文件生成PNG图片。保存在./data/segmentation/labels目录下。
from segmentation.data_util import create_labelsimg_dir = './data/raw/train_imgs/' # 原始图片路径anno_dir = './data/raw/train_labels/labels' # 原始标注文件路径save_dir = './data/segmentation/labels' # 分割标签保存路径create_labels(img_dir, anno_dir, save_dir) # 生成标签print('Finished!')
标签可视化:
import cv2import numpy as npimport matplotlib.pyplot as pltdef show(img_path, label_path): # 读取原始图片 img = cv2.imread(img_path) # 读取标签图片 label = cv2.imread(label_path) # 标注结果图片的背景像素值为0,物体的像素值为1,均显示为黑色。 # 为了可视化效果明显,突出物体的标注区域,这里将物体的像素值1(黑色)替换为255(白色)。 label[label==1] = 255 # 原始图片和标签图片像素融合 merge = img * 0.5 label * 0.5 merge = merge.astype(np.uint8) images_to_observe = [img, label, merge] titles = ['原始图片','标签图片','原始图片 标签图片融合'] # 结果可视化 fig = plt.figure(figsize=(30, 30)) for i in range(len(images_to_observe)): fig.add_subplot(1, len(images_to_observe), i 1).set_title(titles[i], fontsize=18, color='r') imgplot = plt.imshow(images_to_observe[i]) plt.show()if __name__ == '__main__': # 以其中一张图片为例 img_path = './data/raw/train_imgs/train_10.jpg' label_path = './data/segmentation/labels/train_10.png' show(img_path, label_path)
执行以下命令下载预训练模型并保存到./pretrained_model/deeplabv3plus_r50-d8.pth:
!mkdir -p ./pretrained_model!wget --no-check-certificate https://download.openmmlab.com/mmsegmentation/v0.5/deeplabv3plus/deeplabv3plus_r50-d8_512x512_80k_ade20k/deeplabv3plus_r50-d8_512x512_80k_ade20k_20200614_185028-bf1400d8.pth -O ./pretrained_model/deeplabv3plus_r50-d8.pth
训练时加载预训练模型:
./pretrained_model/deeplabv3plus_r50-d8.pth
为了节省训练时间,默认训练1000步,每一步训练8个样本,每200步验证一次精度,每200步保存一次模型。如果要进一步提升精度,可以尝试修改./segmentation/train.py脚本中的配置参数,将训练步数增大。训练好的模型保存在./train_url/segmentation目录下。
训练时默认加载./data/train.txt和./data/eval.txt中保存的样本作为训练集和验证集。请参考./segmentation/train.py。
单卡训练:
单卡大约耗时15分钟左右。
!python ./segmentation/train.py --num_classes=2 --data_url=./data/ --eval_data_url=./data/ --img_dir=raw/train_imgs --ann_dir=segmentation/labels --work_dir=./train_url/segmentation --load_from=./pretrained_model/deeplabv3plus_r50-d8.pth# num_classes:数据类别数# data_url:训练集根目录# eval_data_url:验证集根目录# img_dir:根目录下存放图片的目录名称# ann_dir:根目录下存放标签的目录名称# work_dir:保存输出模型的路径# load_from:预置模型文件路径
多卡训练:
NoteBook创建多卡环境时可使用,用来加速。
!python -m torch.distributed.launch --nproc_per_node=$(/usr/local/nvidia/bin/nvidia-smi -L | wc -l) --master_port=7000 ./segmentation/train.py --launcher=pytorch --num_classes=2 --data_url=./data --eval_data_url=./data --img_dir=raw/train_imgs --ann_dir=segmentation/labels --work_dir=./train_url/segmentation --load_from=./pretrained_model/deeplabv3plus_r50-d8.pth# nproc_per_node:每个节点启动的进程数(GPU数)# launcher:分布式启动方式# num_classes:数据类别数# data_url:训练集根目录# eval_data_url:验证集根目录# img_dir:根目录下存放图片的目录名称# ann_dir:根目录下存放标签的目录名称# work_dir:保存输出模型的路径# load_from:预置模型文件路径
3.1.4.推理 四边形定位 矫正训练完成后,使用训练好的分割模型对图片进行推理,分割模型返回的推理结果是个2维数组,保存了预测图片每个像素点的标签值。 然后对推理结果做四边形的定位,再将四边形矫正成矩形,最后将矩形区域保存成新的图片数据。这部分数据用来训练后面的文本识别算法。保存在./data/textrecog/images_from_seg路径下。
对原始图片./data/raw/train_imgs进行推理:
!python ./segmentation/infer_and_crop.py --num_classes=2 --img_path=./data/raw/train_imgs --checkpoint=./train_url/segmentation/latest.pth --save_dir=./data/textrecog/images_from_seg# num_classes:类别数# img_path:要推理的图片目录# checkpoint:模型文件# save_path:推理结果保存路径
推理结果可视化:
from util import show_pairimg_dir = './data/raw/train_imgs' # 训练集原始图片result_dir = './data/textrecog/images_from_seg' # 训练集推理结果show_pair(img_dir, result_dir, show_num=2)
3.2.文本检测算法与分割算法二选一即可。
3.2.1.数据准备将数据的元信息保存在json文件中。文本检测的标签默认只有一类text,只要对文本区域进行标注并训练,识别出文本区域即可,和分割算法一样,不需要识别文本中的内容。
训练文件保存在./data/textdet/instances_training.json,验证文件保存在./data/textdet/instances_test.json。
import osfrom textdet.data_util import create_water_jsonbase_dir = './data'annotation_dir = os.path.join(base_dir, 'raw/train_labels/labels')img_dir = os.path.join(base_dir, 'raw/train_imgs')# traincreate_water_json(img_dir=img_dir, # 图片路径 anno_dir=annotation_dir, # 标注文件路径 save_file=os.path.join(base_dir, 'textdet/instances_training.json'), # 训练文件保存路径 split=os.path.join(base_dir, 'train.txt')) # 保存训练样本的文件# evalcreate_water_json(img_dir=img_dir, # 图片路径 anno_dir=annotation_dir, # 标注文件路径 save_file=os.path.join(base_dir, 'textdet/instances_test.json'), # 验证文件保存路径 split=os.path.join(base_dir, 'eval.txt')) # 保存验证样本的文件
生成的json文件格式如下:
{ 'images': [{'file_name': 'train_1.jpg', 'height': 960, 'width': 540, 'segm_file': None, 'id': 0}, {'file_name': 'train_2.jpg', 'height': 540, 'width': 960, 'segm_file': None, 'id': 1}, ...]}, 'categories': [{'id': 1, 'name': 'text'}], 'annotations': [{'iscrowd': 0, 'category_id': 1, 'bbox': [126, 283, 241, 147], 'area': 13363, 'segmentation': [[126, 379, 345, 283, 367, 333, 144, 430]], 'image_id': 0, 'id': 0}, {'iscrowd': 0, 'category_id': 1, 'bbox': [430, 159, 204, 67], 'area': 10316, 'segmentation': [[445, 159, 634, 170, 630, 226, 430, 205]], 'image_id': 1, 'id': 1}, ...]}
images:
file_name:图片名称
height: 图片的高
width: 图片的宽
id:图片的索引
categories:
id:标签的索引
name:标签值
annotations:
iscrowd:默认值为0,segmentation使用polygon格式,。
area: 四边形的面积。
segmentation:四边形的四个角点的坐标。
image_id:图片的索引。
id:文本框的索引
!wget --no-check-certificate https://download.pytorch.org/models/resnet50-19c8e357.pth -O ./pretrained_model/resnet50-19c8e357.pth!wget --no-check-certificate https://download.openmmlab.com/mmocr/textdet/dbnet/dbnet_r50dcnv2_fpnc_sbn_2e_synthtext_20210325-aa96e477.pth -O ./pretrained_model/dbnet_r50dcnv2_fpnc_sbn_2e_synthtext.pth
训练的数据集、模型、优化器等均已保存在配置文件./textdet/water_meter_textdet_config.py和./textdet/base_config.py中。base_config.py文件是基础配置,可不必改动。water_meter_textdet_config.py文件是针对该案例的配置,例如数据路径、预训练模型路径、epoch数、学习率等可直接在此文件中修改。
为了节省时间,默认训练10个epoch,每2个epoch验证一次,保存一次模型。如果精度不够,可以修改water_meter_textdet_config.py中的训练参数。将epoch数调大,可调至30-100之间,epoch越大,训练时间越久。训练的模型保存在./train_url/textdet路径下。
单卡训练:
单卡训练较慢,大约耗时15分钟左右。
!python textdet/train.py ./textdet/water_meter_textdet_config.py --work-dir=./train_url/textdet
多卡训练:
NoteBook创建多卡环境时可使用,可加速训练。
!python -m torch.distributed.launch --nproc_per_node=$(/usr/local/nvidia/bin/nvidia-smi -L | wc -l) --master_port=7000 textdet/train.py ./textdet/water_meter_textdet_config.py --launcher=pytorch --work-dir=./train_url/textdet# nproc_per_node:每个节点启动的进程数(GPU数)# launcher:分布式启动方式# work-dir:模型保存的路径
推理生成的结果保存在./data/textdet/results路径下,该路径下包含两个文件夹out_vis_dir和out_txt_dir。out_vis_dir目录下存放每张图片推理的可视化结果,保存为jpg文件。out_txt_dir目录下存放每张图片推理出的四边形区域的坐标值,保存为txt文件。
!python textdet/infer.py ./data/raw/train_imgs ./textdet/water_meter_textdet_config.py ./train_url/textdet/latest.pth --out-dir=./data/textdet/results# 第1个参数: 推理的图片路径# 第2个参数: 模型的配置文件# 第3个参数: 训练好的模型文件# 第4个参数: 推理结果保存路径
out_txt_dir目录下txt文件的内容如下:
393,180,615,177,616,241,394,244,1
393,180,615,177,616,241,394,244:代表4个角点的坐标,分别是左上角,右上角,右下角,左下角。
1:标签索引,1代表是文本区域。
out_vis_dir目录下为图片,可直接进入./data/textdet/results/out_vis_dir目录下打开图片查看,或者使用下面的可视化接口查看:
from util import showimg_dir = './data/textdet/results/out_vis_dir'show(img_dir, show_num=3)
根据上面的推理结果,将文本区域抠出来保存成新的图片,用于训练文本识别模型。保存路径为./data/textrecog/images_from_textdet。
from textdet.data_util import cutoutimg_dir = './data/raw/train_imgs' # 推理时的原始图片路径cutout(img_dir, ann_dir='./data/textdet/results/out_txt_dir/', # 推理结果txt文件 save_dir='./data/textrecog/images_from_textdet' # 抠图结果保存路径 )print('Finished')
抠图结果可视化:
from util import show_pairimg_dir = './data/raw/train_imgs' # 训练集原始图片result_dir = './data/textrecog/images_from_textdet' # 训练集推理结果show_pair(img_dir, result_dir, show_num=2)
本案例数据集物体倾斜角度较小,不需要进行翻转识别,如果遇到随机翻转严重的数据集可以参考这部分代码训练一个识别翻转的模型。
文本检测算法中已经将需要识别的文本区域抠图并保存。接下来利用这部分数据训练一个文本识别算法。因为文本区域做了角度的矫正。但是无法区分翻转0度和180度的图片,而文本识别的算法对翻转180度的场景效果不佳。所以我们又训练了一个专门识别翻转的模型。用来做180度翻转的矫正。
可将文本检测抠出来的图片做180度翻转的离线扩充。根据文本检测抠出的图片以及离线扩充后的数据进行人工标注正常或翻转。
但是本案例图片倾斜角度较小,抠图的结果均为正向,不需要人工标注。原始图片均为正向0度翻转,离线扩充进行180度翻转后的均为180度翻转的图片。
离线扩充:
分割模型的精度较高,结果较准确,我们使用分割模型的推理结果进行离线扩充
import osfrom PIL import Imageimport moxing as moxdef flip_upside_down(img_dir, save_dir): mox.file.make_dirs(save_dir) imgs_list = os.listdir(img_dir) for img_name in imgs_list: img = Image.open(os.path.join(img_dir, img_name)) img.rotate(180).save(os.path.join(save_dir, img_name)) # 图片旋转180度并保存 if __name__ == '__main__': img_dir = './data/textrecog/images_from_seg' flip_save_dir = './data/detect_180/flip_180' # 翻转数据的保存路径,保存在./data/detect_180目录下 flip_upside_down(img_dir, flip_save_dir) # 进行翻转处理 mox.file.copy_parallel(img_dir, './data/detect_180/normal') # 将正常数据拷贝一份到./data/detect_180目录下,用于训练翻转模型 print('Finished!')
离线结果可视化:
from util import showimg_dir = './data/detect_180/flip_180'show(img_dir, show_num=5)
执行以下命令下载模型,模型保存到./pretrained_model/mobilenet_v2.pth。
!wget --no-check-certificate https://download.openmmlab.com/mmclassification/v0/mobilenet_v2/mobilenet_v2_batch256_imagenet_20200708-3b2dc3af.pth -O ./pretrained_model/mobilenet_v2.pth
识别图片翻转可以使用小点的模型,这里我们以mobilenet_v2为例,如果精度不够可以换稍大点的模型。
使用前面下载的预训练模型,./pretrained_model/mobilenet_v2.pth。
为了节省时间,默认训练25个epoch,每一步训练64个样本,每2个epoch验证一次精度,每2个epoch保存一次模型。大约耗时2-3分钟左右。如果想要更高的精度,可以修改./code/detect_180/train.py脚本中的配置参数。或者换大点的模型。模型较小,使用单卡训练即可。多卡训练可参考其他模型的多卡训练。
这里没有将数据切分成训练集和验证集,可自行切分,指定eval_data_url。
!python ./detect_180/train.py --num_classes=2 --data_url=./data/detect_180 --work_dir=./train_url/detect_180 --load_from=./pretrained_model/mobilenet_v2.pth#--eval_data_url=./data/detect_180# num_classes:数据类别数# data_url:训练集存放路径# eval_data_url:验证集存放路径# work_dir:保存输出模型的路径# load_from:预置模型文件路径
训练的模型保存在./train_url/detect_180路径下,加载该路径下的模型进行推理。就以训练集路径下的翻转图片为例。
!python ./detect_180/infer.py --num_classes=2 --img_path=./data/detect_180/flip_180 --checkpoint=./train_url/detect_180/latest.pth
根据文本检测后的抠图数据和原始的txt标注文件创建文本内容识别的标签文件。
水表中有的度数转到一半,会同时出现两个数字。原始标注文件中也标注了多个label,从实际情况出发,我们选择符合实际情况的标签(如果转到一半,选择小的那个度数)。有一些特殊的图片,需要人工过滤一下:例如train_19.jpg这张图片。
import cv2import matplotlib.pyplot as plt imgplot = plt.imshow(cv2.imread('./data/raw/train_imgs/train_19.jpg'))plt.show()
该图片的原始标签有四个['00069', '00070', '00060', '00079'],但从实际情况出发,符合条件的应该是00069和 00070。我们选择度数小的那个作为标签值,即00069。
还有几张图片的标注格式与其他标注文件不同,例如train_140.jpg这张图片标注的标签格式是00470.00479,其他文件都是空格分隔00470 00479。这种图片只有几张我们也进行人工过滤。
下面根据原始的标注文件生成文本识别需要的标注文件格式。
from textrecog.data_util import gen_label_fileann_dir = './data/raw/train_labels/labels'# traintrain_split_path = './data/train.txt'train_save_path = './data/textrecog/train_label.txt'gen_label_file(train_split_path, ann_dir, train_save_path)# evaleval_split_path = './data/eval.txt'eval_save_path = './data/textrecog/eval_label.txt'gen_label_file(eval_split_path, ann_dir, eval_save_path)print('Finished!')print('标注文件内容如下,文本图片名称 文本数字内容')!head -n 5 ./data/textrecog/train_label.txt
执行以下命令下载模型,模型保存到./pretrained_model/crnn_academic-a723a1c5.pth。
!wget --no-check-certificate https://download.openmmlab.com/mmocr/textrecog/crnn/crnn_academic-a723a1c5.pth -O ./pretrained_model/crnn_academic-a723a1c5.pth
训练的数据集、模型、优化器等均已保存在配置文件./textrecog/water_meter_textrecog_config.py和./textrecog/base_config.py中。base_config.py文件是基础配置,可不必改动。water_meter_textrecog_config.py文件是针对该案例的配置,例如数据路径、预训练模型路径、epoch数、学习率等可直接在此文件中修改。
默认训练50个epoch,精度不够,可以修改water_meter_textrecog_config.py中的训练参数进行调参。
训练生成的模型保存在./train_url/textrecog路径下。大约耗时2-3分钟左右:
!python ./textrecog/train.py ./textrecog/water_meter_textrecog_config.py --work-dir=./train_url/textrecog
生成推理文件列表:
将需要推理的文件名称保存在txt文件中。
import osimg_root_path = './data/textrecog/images_from_seg'with open('./data/textrecog/infer_imgs_list.txt', 'w') as f: f.write(''.join(os.listdir(img_root_path)))!head -n 5 ./data/textrecog/infer_imgs_list.txt # 查看生成文件的前5行样本
推理:
利用训练好的模型进行推理并保存推理结果。结果保存在./data/textrecog/infer_results目录下。该目录下会有多个文件和目录生成。我们只需查看out_vis_dir和result.txt即可。
!mkdir ./data/textrecog/infer_results # 创建推理结果的保存路径!python ./textrecog/infer.py ./data/textrecog/images_from_seg ./data/textrecog/infer_imgs_list.txt ./textrecog/water_meter_textrecog_config.py ./train_url/textrecog/latest.pth --out_dir=./data/textrecog/infer_results# 第1个参数: 推理的图片根目录# 第2个参数: 图片根目录下的推理文件列表# 第3个参数: 模型配置文件# 第4个参数: 训练好的模型文件# 第5个参数: 推理结果保存路径
推理结果可视化:
from util import showimg_dir = './data/textrecog/infer_results/out_vis_dir'show(img_dir, show_num=5)
以上我们已经完成了水表识别的整个流程。现在我们将上述算法串联在一起,实现端到端的预测。输入数据为原始拍摄的图片。输出为水表的度数。
与文本识别算法算法二选一。
如果文本检测部分使用的是分割算法,参考infer_end2end_with_seg.py脚本,模型路径,数据路径都配置在了该脚本中,可修改该脚本。
如果前面的模型或者数据有改动部分,该脚本请同步修改。
!python infer_end2end_with_seg.py --num_classes=2 --img_path=./data/raw/test_imgs --out=./outputs_with_seg --detect_180=False# num_classes:类别数# img_path:推理图片路径# out:推理结果保存路径# detect_180:如果训练了识别翻转的模型,设置detect_180=True即可,默认值为False
推理结果可视化:
from util import show_pairimg_dir = './data/raw/test_imgs'res_dir = './outputs_with_seg'show_pair(img_dir, res_dir, show_num=2)
与分割算法算法二选一。
如果文本检测部分使用的是文本识别算法,参考infer_end2end_with_textdet.py脚本,模型路径,数据路径都配置在了该脚本中,可修改该脚本。
如果前面的模型或者数据有改动部分,该脚本请同步修改。
!python infer_end2end_with_textdet.py --num_classes=2 --img_path=./data/raw/test_imgs --out=./outputs_with_textdet --detect_180=False# num_classes:类别数# img_path:推理图片路径# out:推理结果保存路径# detect_180:如果训练了识别翻转的模型,设置detect_180=True即可,默认值为False
推理结果可视化:
from util import show_pairimg_dir = './data/raw/test_imgs'res_dir = './outputs_with_textdet'show_pair(img_dir, res_dir, show_num=2)
花粉社群VIP加油站
猜你喜欢