基于MMOCR的车牌识别
Created: June 8, 2023 4:41 PM
项目类型: 学习项目
状态: 已完成
使用方法
安装
首先自行安装pytorch
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu116
然后使用官网提供的命令安装
pip install -U openmim
mim install mmengine
mim install mmcv
mim install mmdet
然后用源码的方式安装MMOCR,这样方便开发,否则只能当成一个依赖库来使用
git clone https://github.com/open-mmlab/mmocr.git
cd mmocr
pip install -v -e .
⚠️ 在Linux中,用mim install mmengine
或者pip install mmengine
方法安装的mmengine最高版本是0.7.0,但我们可以看到,至少需要0.7.1。从下面的报错可以看出来。
root@autodl-container-acda11befa-0322c85c:~/mmocr# mim install mmengine==0.7.1 Looking in indexes: [https://repo.huaweicloud.com/repository/pypi/simple](https://repo.huaweicloud.com/repository/pypi/simple) Looking in links: [https://download.openmmlab.com/mmcv/dist/cu113/torch1.10.0/index.html](https://download.openmmlab.com/mmcv/dist/cu113/torch1.10.0/index.html) ERROR: Could not find a version that satisfies the requirement mmengine==0.7.1 (from versions: 0.0.1rc0, 0.1.0, 0.2.0, 0.3.0, 0.3.1, 0.3.2, 0.4.0, 0.5.0, 0.6.0, 0.7.0) ERROR: No matching distribution found for mmengine==0.7.1
MMOCR | MMEngine | MMCV | MMDetection | ||
---|---|---|---|---|---|
dev-1.x | 0.7.1 <= mmengine < 1.0.0 | 2.0.0rc4 <= mmcv < 2.1.0 | 3.0.0rc5 <= mmdet < 3.1.0 | ||
1.0.0 | 0.7.1 <= mmengine < 1.0.0 | 2.0.0rc4 <= mmcv < 2.1.0 | 3.0.0rc5 <= mmdet < 3.1.0 | ||
1.0.0rc6 | 0.6.0 <= mmengine < 1.0.0 | 2.0.0rc4 <= mmcv < 2.1.0 | 3.0.0rc5 <= mmdet < 3.1.0 | ||
1.0.0rc[4-5] | 0.1.0 <= mmengine < 1.0.0 | 2.0.0rc1 <= mmcv < 2.1.0 | 3.0.0rc0 <= mmdet < 3.1.0 | ||
1.0.0rc[0-3] | 0.0.0 <= mmengine < 0.2.0 | 2.0.0rc1 <= mmcv < 2.1.0 | 3.0.0rc0 <= mmdet < 3.1.0 |
但是在清华源有0.7.3版本的mmengine,可以用这条命令来下载pip install mmengine==0.7.3 -i [https://pypi.tuna.tsinghua.edu.cn/simple](https://pypi.tuna.tsinghua.edu.cn/simple)
这条callout总结于2023.06.09,可能以后会修好这个bug,并且此时我在AutoDL里存好可以用的MMOCR镜像了。
安装好后可以在python中运行如下命令检验
>>> from mmocr.apis import MMOCRInferencer
>>> ocr = MMOCRInferencer(det='DBNet', rec='CRNN')
>> ocr('demo/demo_text_ocr.jpg', show=True, print_result=True)
若 MMOCR 的安装无误,你在这一节完成后应当能看到以图片和文字形式表示的识别结果:
# 识别结果{'predictions':[{'rec_texts':['cbanks', 'docecea', 'grouf', 'pwate', 'chobnsonsg', 'soxee', 'oeioh', 'c', 'sones', 'lbrandec', 'sretalg', '11', 'to8', 'round', 'sale', 'year',
'ally', 'sie', 'sall'], 'rec_scores':[...], 'det_polygons':[...], 'det_scores':
[...]}]}
使用
使用前需要需要**from** **mmocr.apis** **import** **MMOCRInferencer**
两种推理器
MMOCR 中存在两种不同的推理器:
- 标准推理器:MMOCR 中的每个基本任务都有一个标准推理器,即
TextDetInferencer
(文本检测),TextRecInferencer
(文本识别),TextSpottingInferencer
(端到端 OCR) 和KIEInferencer
(关键信息提取)。它们具有非常相似的接口,具有标准的输入/输出协议,并且总体遵循 OpenMMLab 的设计。这些推理器也可以被串联在一起,以便对一系列任务进行推理。- MMOCRInferencer:我们还提供了
MMOCRInferencer
,一个专门为 MMOCR 设计的便捷推理接口。它封装和链接了 MMOCR 中的所有推理器,因此用户可以使用此推理器对图像执行一系列任务,并直接获得最终结果。但是,它的接口与标准推理器有一些不同,并且为了简单起见,可能会牺牲一些标准的推理器功能。
我暂时只需要使用MMOCRInferencer
需要注意推理器的返回值,这是手册对输出结果的介绍:
{
'predictions' : [
# 每个实例都对应于一个输入图像
{
'det_polygons': [...], # 2d 列表,长度为 (N,),格式为 [x1, y1, x2, y2, ...]
'det_scores': [...], # 浮点列表,长度为(N, )
'det_bboxes': [...], # 2d 列表,形状为 (N, 4),格式为 [min_x, min_y, max_x, max_y]
'rec_texts': [...], # 字符串列表,长度为(N, )
'rec_scores': [...], # 浮点列表,长度为(N, )
'kie_labels': [...], # 节点标签,长度为 (N, )
'kie_scores': [...], # 节点置信度,长度为 (N, )
'kie_edge_scores': [...], # 边预测置信度, 形状为 (N, N)
'kie_edge_labels': [...] # 边标签, 形状为 (N, N)
},
...
],
'visualization' : [
array(..., dtype=uint8),
]
}
detector = MMOCRInferencer(det='TextSnake', # 文本检测算法,这里指定为 TextSnake,也可替换为其他 MMOCR 支持的文本区域检测算法
rec=None, # 文本识别算法,这里不指定,也可替换为其他 MMOCR 支持的文本识别算法
device='cuda') # 指定运算设备为 cpu 或 cuda
result = detector('demo/demo_text_det.jpg', out_dir='outputs/demo_text_det_pred.jpg',save_vis='True')
在这记录一个引用例子,因为手册没有给result的引用例子
len(result['predictions'][0]['det_polygons'])
结果是9,说明识别到9个文字区域。
'det_polygons'
是文字的多边形节点的横纵坐标,'det_scores'
为置信度
数据集配置
CCPD
CCPD的数据集其实在 https://opendatalab.com/ 中可以下载,可惜我一开始不知道,所以自己配置的数据集。
CCPD是一个中国车牌数据集,下图为其中一个示例。
这个数据集的标注就在文件名字上,其规则为:
图片命名:“025-95_113-154&383_386&473-386&473_177&454_154&383_363&402-0_0_22_27_27_33_16-37-15.jpg”
解释:
1、025:车牌区域占整个画面的比例;
2、95_113: 车牌水平和垂直角度, 水平95°, 竖直113°
3、154&383_386&473:标注框左上、右下坐标,左上(154, 383), 右下(386, 473)
4、86&473_177&454_154&383_363&402:标注框四个角点坐标,顺序为右下、左下、左上、右上
5、0_0_22_27_27_33_16:车牌号码映射关系如下: 第一个0为省份 对应省份字典provinces中的’皖’,;第二个0是该车所在地的地市一级代码,对应地市一级代码字典alphabets的’A’;后5位为字母和文字, 查看车牌号ads字典,如22为Y,27为3,33为9,16为S,最终车牌号码为皖AY339S6、64为亮度,数值越大车牌越亮(可能不准确,仅供参考);
7、233为模糊度,数值越小车牌越模糊(可能不准确,仅供参考)。
# 每个数组的最后一个字符是字母O而不是数字0。我们使用O作为“无字符”的标志,因为中文车牌字符中没有O。 provinces = ["皖", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", "苏", "浙", "京", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤", "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新", "警", "学", "O"] alphabets = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'O'] ads = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'O']
MMOCR标注格式
MMOCR有自己的标准标注格式,我们需要转换成那种格式。首先我们看一下格式的样子,格式来源:https://mmocr.readthedocs.io/zh_CN/dev-1.x/basic_concepts/datasets.html#id4
此标注文件是一个
.json
文件,存储一个包含metainfo
和data_list
的dict
,前者包括有关数据集的基本信息,后者由每个图片的标注组成。这里呈现了标注文件中的所有字段的列表,但其中某些字段仅会在特定任务中被用到。
{
"metainfo":
{
"dataset_type": "TextDetDataset", # 可选项: TextDetDataset/TextRecogDataset/TextSpotterDataset
"task_name": "textdet", # 可选项: textdet/textspotter/textrecog
"category": [{"id": 0, "name": "text"}] # 在 textdet/textspotter 里用到
},
"data_list":
[
{
"img_path": "test_img.jpg", #这里要写出具体的文件相对位置,而不是仅仅写文件名。
"height": 604,
"width": 640,
"instances": # 一图内的多个实例
[
{
"bbox": [0, 0, 10, 20], # textdet/textspotter 内用到, [x1, y1, x2, y2]。
"bbox_label": 0, # 对象类别, 在 MMOCR 中恒为 0 (文本)
"polygon": [0, 0, 0, 10, 10, 20, 20, 0], # textdet/textspotter 内用到。 [x1, y1, x2, y2, ....]
"text": "mmocr", # textspotter/textrecog 内用到
"ignore": False # textspotter/textdet 内用到,决定是否在训练时忽略该实例
},
#...
],
}
#... 多图片
]
}
观察格式我们可以发现,我们需要的标注只有"bbox""polygon""text"
CCPD其他信息都用不上。
现在我们就要把这些需要的信息整理成标准格式,下面是读取文件名,将其整理成json格式的代码
import json
import os
provinces = ["皖", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", "苏", "浙", "京", "闽", "赣", "鲁", "豫", "鄂",
"湘", "粤", "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新", "警", "学", "O"]
alphabets = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W',
'X', 'Y', 'Z', 'O']
ads = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X',
'Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'O']
data_list = []
invalid = []
folder_path = r"E:\dataset\CCPDdet\textdet_imgs\test"
pic_num = 0
# print(folder_path.split('\\')[-2:]) 测试返回的是不是:textdet_imgs\test,发现不是,返回的是textdet_imgs和test,是分开的
# 遍历整个文件夹
for fi in os.listdir(folder_path):
if fi.endswith(".jpg"): # 文件是以jpg结尾的
text = ''
# 图片名字的每个部分是用-分割的,6是最多分割次数,也就是最多会分割成7份
file_name = fi.split("-", 6)
# 有些图片命名不正确,筛选出来删掉
if len(file_name) < 7:
os.remove(os.path.join(folder_path, fi))
invalid.append(fi)
continue
# 提取bbox
plate_bbox = file_name[2].split("_", 1)
x1 = int(plate_bbox[0].split("&", 1)[0])
y1 = int(plate_bbox[0].split("&", 1)[1])
x2 = int(plate_bbox[1].split("&", 1)[0])
y2 = int(plate_bbox[1].split("&", 1)[1])
pl_bb = [x1, y1, x2, y2]
# 提取polygon
plate_polygon = file_name[3].split("_", 3)
px1 = int(plate_polygon[0].split("&", 1)[0])
py1 = int(plate_polygon[0].split("&", 1)[1])
px2 = int(plate_polygon[1].split("&", 1)[0])
py2 = int(plate_polygon[1].split("&", 1)[1])
px3 = int(plate_polygon[2].split("&", 1)[0])
py3 = int(plate_polygon[2].split("&", 1)[1])
px4 = int(plate_polygon[3].split("&", 1)[0])
py4 = int(plate_polygon[3].split("&", 1)[1])
pl_po = [px1, py1, px2, py2, px3, py3, px4, py4]
palte_text = file_name[4].split("_", 7)
text_len = len(palte_text)
text += provinces[int(palte_text[0])]
text += alphabets[int(palte_text[1])]
text += ads[int(palte_text[2])] + ads[int(palte_text[3])] + ads[int(palte_text[4])] + ads[
int(palte_text[5])] + ads[int(palte_text[6])]
# 新能源的车牌是多一位的,需要多读取一位
if text_len == 8:
text += ads[int(palte_text[7])]
new_instance = {
# img_path是数据集根目录到图片之间的相对地址的,这里巨坑,手册也不说清楚,还好问题不大,看报错就知道了
"img_path": os.path.join(os.path.join(folder_path.split('\\')[-2], folder_path.split('\\')[-1]), fi),
"height": 1160,
"width": 720,
"instances":
[
{
"bbox": pl_bb,
"bbox_label": 0,
"polygon": pl_po,
"text": text,
"ignore": False
},
],
}
data_list.append(new_instance)
print(invalid)
print(pic_num)
print(len(data_list))
json_data = {
"metainfo": {
"dataset_type": "TextDetDataset",
"task_name": "textdet",
"category": [{"id": 0, "name": "text"}]
},
"data_list": data_list
}
# 一定要encoding="utf-8",因为有中文,不这么弄会报个和utf-8有关的错误
with open("textdet_test.json", "w", encoding="utf-8") as f:
json.dump(json_data, f, indent=4, ensure_ascii=False)
以下是标记好后的一个范例
{
"metainfo": {
"dataset_type": "TextDetDataset",
"task_name": "textdet",
"category": [
{
"id": 0,
"name": "text"
}
]
},
"data_list": [
{
"img_path": "textdet_imgs\\test\\0018-3_1-284&522_343&548-341&548_284&545_286&522_343&525-0_0_18_31_29_32_8-74-14.jpg",
"height": 1160,
"width": 720,
"instances": [
{
"bbox": [
284,
522,
343,
548
],
"bbox_label": 0,
"polygon": [
341,
548,
284,
545,
286,
522,
343,
525
],
"text": "皖AU758J",
"ignore": false
}
]
},
...
配置完json,我们要看下图片和json是怎么放置的。在 https://mmocr.readthedocs.io/zh_CN/dev-1.x/user_guides/dataset_prepare.html#id4 中说到MMOCR格式:
data/icdar2015
├── textdet_imgs
│ ├── test
│ └── train
├── textdet_test.json
└── textdet_train.json
我感觉这个格式不重要,而且也标注格式的image_path有点冲突,导致image_path要写textdet_imgs\test\xxx.jpg或者textdet_imgs\train’\xxx.jpg。
为了弄出这样的格式,要用python的脚本来搬运图像,因为CCPD文件夹下有很多子文件夹,子文件夹里才是图像。手动搬运比较麻烦,以下是搬运用的代码。
import shutil
import os
import glob
import random
folder_path = r"E:\数据集\CCPD"
test_path = r"E:\数据集\CCPDdet\textdet_imgs\test"
train_path = r"E:\数据集\CCPDdet\textdet_imgs\train"
# 用嵌套的方法进入每一个子文件夹
def getFileList(dir, Filelist, ext=None):
"""
获取文件夹及其子文件夹中文件列表
输入 dir:文件夹根目录
输入 ext: 扩展名
返回: 文件路径列表
"""
global test_path
global train_path
if os.path.isfile(dir): # 如果文件路径是文件,这执行文件的判断,看是否是所需类型
if ext is None:
Filelist.append(dir)
else:
if dir.endswith(ext): # 检查文件后缀,看是否是满足要求的文件类型
if (random.randint(1, 10)) < 8: # 划分数据集和测试集
# 判断文件是否存在,CCPD数据集会有重复的图片放在不同的文件夹里
if not (os.path.exists(os.path.join(train_path,dir.split('\\')[-1]))): # 在Linux中,需要改成split('/')
shutil.move(dir, train_path)
else:
if not (os.path.exists(os.path.join(test_path, dir.split('\\')[-1]))):
shutil.move(dir, test_path)
Filelist.append(dir)
elif os.path.isdir(dir): # 如果文件路径是文件夹,则执行嵌套,继续访问子文件夹
for s in os.listdir(dir):
newDir = os.path.join(dir, s)
getFileList(newDir, Filelist, ext)
return Filelist
Filelist = getFileList(folder_path, [], ".jpg")
print(len(Filelist))
然后就是数据集config的配置,这里有 https://mmocr.readthedocs.io/zh_CN/dev-1.x/user_guides/dataset_prepare.html 示例
icdar2015_textdet_data_root = 'data/icdar2015' # 数据集根目录
# 训练集配置
icdar2015_textdet_train = dict(
type='OCRDataset',
data_root=icdar2015_textdet_data_root, # 数据根目录
ann_file='textdet_train.json', # 标注文件名称
filter_cfg=dict(filter_empty_gt=True, min_size=32), # 数据过滤
pipeline=None)
# 测试集配置
icdar2015_textdet_test = dict(
type='OCRDataset',
data_root=icdar2015_textdet_data_root,
ann_file='textdet_test.json',
test_mode=True,
pipeline=None)
这个示例文件在mmocr\configs\textdet\*base*\datasets\[icdar2015.py](http://icdar2015.py/)
中可以找到。然后我修改成CCPD需要的样式。这里的ccpd_textdet_data_root
指定了数据集根目录,所以上面json标注的image_path指的是相对于数据集根目录的相对路径。也许可以尝试更改字典里的data_root
解决,不过这次没有尝试。
ccpd_textdet_data_root = 'E:\dataset\CCPDdet'
ccpd_textdet_train = dict(
type='OCRDataset',
data_root=ccpd_textdet_data_root,
ann_file='textdet_train.json',
filter_cfg=dict(filter_empty_gt=True, min_size=32),
pipeline=None)
ccpd_textdet_test = dict(
type='OCRDataset',
data_root=ccpd_textdet_data_root,
ann_file='textdet_test.json',
test_mode=True,
pipeline=None)
接下来要检验标注是否正确,可以使用可视化工具。在 https://mmocr.readthedocs.io/zh_CN/dev-1.x/user_guides/dataset_prepare.html 中提到一个工具 browse_dataset.py
可以通过这两个链接去看看这工具。下面是可视化的一个例子。
数据准备完毕以后,你也可以通过使用我们提供的数据集浏览工具 browse_dataset.py 来可视化数据集的标签是否被正确生成,例如:
python tools/analysis_tools/browse_dataset.py configs/textdet/_base_/datasets/icdar2015.py
参数名 | 类型 | 描述 |
---|---|---|
config | str | (必须) 配置文件路径。 |
-o, --output-dir | str | 如果图形化界面不可用,请指定一个输出路径来保存可视化结果。 |
-p, --phase | str | 用于指定需要可视化的数据集切片,如 “train”, “test”, “val”。当数据集存在多个变种时,也可以通过该参数来指定待可视化的切片。 |
-m, --mode | original, transformed, pipeline | 用于指定数据可视化的模式。original:原始模式,仅可视化数据集的原始标注;transformed:变换模式,展示经过所有数据变换步骤的最终图像;pipeline:流水线模式,展示数据变换过程中每一个中间步骤的变换图像。默认使用 transformed 变换模式。 |
-t, --task | auto, textdet, textrecog | 用于指定可视化数据集的任务类型。auto:自动模式,将依据给定的配置文件自动选择合适的任务类型,如果无法自动获取任务类型,则需要用户手动指定为 textdet 文本检测任务 或 textrecog 文本识别任务。默认采用 auto 自动模式。 |
-n, --show-number | int | 指定需要可视化的样本数量。若该参数缺省则默认将可视化全部图片。 |
-i, --show-interval | float | 可视化图像间隔时间,默认为 2 秒。 |
–cfg-options | float | 用于覆盖配置文件中的参数,详见https://mmocr.readthedocs.io/zh_CN/dev-1.x/user_guides/config.md#command-line-modification。 |
# python裁剪图片并保存
from PIL import Image
import os
import cv2
srcPath = r"E:\dataset\CCPDdet\textdet_imgs\train"
dstPath = r"E:\dataset\CCPDrec\train"
# 读取图片
for fi in os.listdir(srcPath):
if fi.endswith(".jpg"):
file_name = fi.split("-", 6)
plate_bbox = file_name[2].split("_", 1)
x1 = int(plate_bbox[0].split("&", 1)[0])
y1 = int(plate_bbox[0].split("&", 1)[1])
x2 = int(plate_bbox[1].split("&", 1)[0])
y2 = int(plate_bbox[1].split("&", 1)[1])
pl_bb = (x1, y1, x2, y2)
img = Image.open(os.path.join(srcPath, fi))
img_rec = img.crop(pl_bb)
savs_path = os.path.join(dstPath, fi)
img_rec.save(savs_path)
训练
文本检测训练
训练的配置
在 https://mmocr.readthedocs.io/zh_CN/latest/user_guides/config.html 中有更为详细的配置
上述网址提到的目录结构
目录结构
MMOCR
所有配置文件都放置在configs
文件夹下。为了避免配置文件过长,同时提高配置文件的可复用性以及清晰性,MMOCR 利用 Config 文件的继承特性,将配置内容的八个部分做了拆分。因为每部分均与算法任务相关,因此 MMOCR 对每个任务在 Config 中提供了一个任务文件夹,即textdet
(文字检测任务)、textrecog
(文字识别任务)、kie
(关键信息提取)。同时各个任务算法配置文件夹下进一步划分为两个部分:_base_
文件夹与诸多算法文件夹:
_base_
文件夹下主要存放与具体算法无关的一些通用配置文件,各部分依目录分为常用的数据集、常用的训练策略以及通用的运行配置。- 算法配置文件夹中存放与算法强相关的配置项。算法配置文件夹主要分为两部分:算法的模型与数据流水线:OCR 领域中一般情况下数据增强策略与算法强相关,因此模型与数据流水线通常置于统一位置。算法在制定数据集上的特定配置:用于训练和测试的配置,将分散在不同位置的 base 配置汇总。同时可能会修改一些
_base_
中的变量,如batch size, 数据流水线,训练策略等最后的将配置内容中的各个模块分布在不同配置文件中,最终各配置文件内容如下:
最终目录结构如下:
configs
├── textdet
│ ├── _base_
│ │ ├── datasets
│ │ │ ├── icdar2015.py
│ │ │ ├── icdar2017.py
│ │ │ └── totaltext.py
│ │ ├── schedules
│ │ │ └── schedule_adam_600e.py
│ │ └── default_runtime.py
│ └── dbnet
│ ├── _base_dbnet_resnet18_fpnc.py
│ └── dbnet_resnet18_fpnc_1200e_icdar2015.py
├── textrecog
│ ├── _base_
│ │ ├── datasets
│ │ │ ├── icdar2015.py
│ │ │ ├── icdar2017.py
│ │ │ └── totaltext.py
│ │ ├── schedules
│ │ │ └── schedule_adam_base.py
│ │ └── default_runtime.py
│ └── crnn
│ ├── _base_crnn_mini-vgg.py
│ └── crnn_mini-vgg_5e_mj.py
└── kie
├── _base_
│ ├──datasets
│ └── default_runtime.py
└── sgdmr
└── sdmgr_novisual_60e_wildreceipt_openset.py
针对以上信息,提出一个简易的配置流程:
1、配置数据集的.py,这一步在前面讲述过了,不再重复
2、配置网络的的.py,对应上面的dbnet_resnet18_fpnc_1200e_icdar2015.py。下面是我配置好的CCPD,这里需要注意batch_size=4和num_workers。在训练时有一点需要注意,训练时可以开启混合精度,开启后可以将batch_size设置大一些,但好像在训练几个epoch进行val(在测试集上测试,用英文是为了与手动测试区别开)的时候,混合精度就失效了?总之会导致爆内存,需要将val的batch_size设置小一些。这个文件在开头引用了'_base_dbnetpp_resnet50-dcnv2_fpnc.py'
,这个文件决定了训练用什么网络。在 https://mmocr.readthedocs.io/zh_CN/dev-1.x/modelzoo.html#id2 可以看到不同的网络,我选了精度比较高的**DBNetpp
。**
_base_ = [
'_base_dbnetpp_resnet50-dcnv2_fpnc.py',
'../_base_/default_runtime.py',
'../_base_/datasets/ccpd.py',
'../_base_/schedules/schedule_sgd_1200e.py',
]
load_from = 'https://download.openmmlab.com/mmocr/textdet/dbnetpp/tmp_1.0_pretrain/dbnetpp_r50dcnv2_fpnc_100k_iter_synthtext-20220502-352fec8a.pth' # noqa
# dataset settings
train_list = [_base_.ccpd_textdet_train]
test_list = [_base_.ccpd_textdet_test]
train_dataloader = dict(
batch_size=4,
num_workers=1,
persistent_workers=True,
sampler=dict(type='DefaultSampler', shuffle=True),
dataset=dict(
type='ConcatDataset',
datasets=train_list,
pipeline=_base_.train_pipeline))
val_dataloader = dict(
batch_size=4,
num_workers=1,
persistent_workers=True,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=dict(
type='ConcatDataset',
datasets=test_list,
pipeline=_base_.test_pipeline))
test_dataloader = val_dataloader
auto_scale_lr = dict(base_batch_size=16)
3、配置训练略的.py,对应上面的schedule_adam_600e.py,下面是配置好的CCPD。这里要注意训练epochs,学习率及其调整策略。
# optimizer
optim_wrapper = dict(
type='OptimWrapper',
optimizer=dict(type='SGD', lr=0.007, momentum=0.9, weight_decay=0.0001))
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=10, val_interval=2)
val_cfg = dict(type='ValLoop')
test_cfg = dict(type='TestLoop')
# learning policy
param_scheduler = dict(type='CosineAnnealingLR', by_epoch=True, T_max=10)
4、配置对应上面的default_runtime.py。这里需要注意checkpoint,调整多久保存一次权重。
default_scope = 'mmocr'
env_cfg = dict(
cudnn_benchmark=False,
mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
dist_cfg=dict(backend='nccl'),
)
randomness = dict(seed=None)
default_hooks = dict(
timer=dict(type='IterTimerHook'),
logger=dict(type='LoggerHook', interval=1),
param_scheduler=dict(type='ParamSchedulerHook'),
checkpoint=dict(type='CheckpointHook', interval=1),
sampler_seed=dict(type='DistSamplerSeedHook'),
sync_buffer=dict(type='SyncBuffersHook'),
visualization=dict(
type='VisualizationHook',
interval=1,
enable=False,
show=False,
draw_gt=False,
draw_pred=False),
)
# Logging
log_level = 'INFO'
log_processor = dict(type='LogProcessor', window_size=10, by_epoch=True)
load_from = None
resume = False
# Evaluation
val_evaluator = dict(type='HmeanIOUMetric')
test_evaluator = val_evaluator
# Visualization
vis_backends = [dict(type='LocalVisBackend')]
visualizer = dict(
type='TextDetLocalVisualizer',
name='visualizer',
vis_backends=vis_backends)
配置文件里能配置的东西很多,具体要看官网文档里的细节。配置完就可以开始训练了。
开始训练
训练的文档在 https://mmocr.readthedocs.io/zh_CN/latest/user_guides/train_test.html,下面是文档提到的训练代码
# 通过调用 tools/train.py 来训练指定的 MMOCR 模型CUDA_VISIBLE_DEVICES= python tools/train.py${CONFIG_FILE}[PY_ARGS]
# 训练
# 示例 1:使用 CPU 训练 DBNet
CUDA_VISIBLE_DEVICES=-1 python tools/train.py configs/textdet/dbnet/dbnet_resnet50-dcnv2_fpnc_1200e_icdar2015.py
# 示例 2:指定使用 gpu:0 训练 DBNet,指定工作目录为 dbnet/,并打开混合精度(amp)训练
CUDA_VISIBLE_DEVICES=0 python tools/train.py configs/textdet/dbnet/dbnet_resnet50-dcnv2_fpnc_1200e_icdar2015.py --work-dir dbnet/ --amp
下表列出了 train.py
支持的所有参数。其中,不带 --
前缀的参数为必须的位置参数,带 --
前缀的参数为可选参数
参数 | 类型 | 说明 |
---|---|---|
config | str | (必须)配置文件路径。 |
–work-dir | str | 指定工作目录,用于存放训练日志以及模型 checkpoints。 |
–resume | bool | 是否从断点处恢复训练。 |
–amp | bool | 是否使用混合精度。 |
–auto-scale-lr | bool | 是否使用学习率自动缩放。 |
–cfg-options | str | 用于覆写配置文件中的指定参数。https://mmocr.readthedocs.io/zh_CN/latest/user_guides/train_test.html#%E6%B7%BB%E5%8A%A0%E7%A4%BA%E4%BE%8B |
–launcher | str | 启动器选项,可选项目为 [‘none’, ‘pytorch’, ‘slurm’, ‘mpi’]。 |
–local_rank | int | 本地机器编号,用于多机多卡分布式训练,默认为 0。 |
文本内容识别训练
训练配置
训练配置的内容基本和文本检测差不多,只是有一些其他内容不太一样。文本内容识别网络选择的是**ABINet,**其原始的配置文件abinet_20e_st-an_mj.py
如下。
_base_ = [
'../_base_/datasets/mjsynth.py',
'../_base_/datasets/synthtext.py',
'../_base_/datasets/cute80.py',
'../_base_/datasets/iiit5k.py',
'../_base_/datasets/svt.py',
'../_base_/datasets/svtp.py',
'../_base_/datasets/icdar2013.py',
'../_base_/datasets/icdar2015.py',
'../_base_/default_runtime.py',
'../_base_/schedules/schedule_adam_base.py',
'_base_abinet.py',
]
load_from = 'https://download.openmmlab.com/mmocr/textrecog/abinet/abinet_pretrain-45deac15.pth' # noqa
optim_wrapper = dict(optimizer=dict(lr=1e-4))
train_cfg = dict(max_epochs=20)
# learning policy
param_scheduler = [
dict(
type='LinearLR', end=2, start_factor=0.001,
convert_to_iter_based=True),
dict(type='MultiStepLR', milestones=[16, 18], end=20),
]
# dataset settings
train_list = [
_base_.mjsynth_textrecog_train, _base_.synthtext_an_textrecog_train
]
test_list = [
_base_.cute80_textrecog_test, _base_.iiit5k_textrecog_test,
_base_.svt_textrecog_test, _base_.svtp_textrecog_test,
_base_.icdar2013_textrecog_test, _base_.icdar2015_textrecog_test
]
train_dataset = dict(
type='ConcatDataset', datasets=train_list, pipeline=_base_.train_pipeline)
test_dataset = dict(
type='ConcatDataset', datasets=test_list, pipeline=_base_.test_pipeline)
train_dataloader = dict(
batch_size=192,
num_workers=32,
persistent_workers=True,
sampler=dict(type='DefaultSampler', shuffle=True),
dataset=train_dataset)
test_dataloader = dict(
batch_size=1,
num_workers=4,
persistent_workers=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=test_dataset)
val_dataloader = test_dataloader
val_evaluator = dict(
dataset_prefixes=['CUTE80', 'IIIT5K', 'SVT', 'SVTP', 'IC13', 'IC15'])
test_evaluator = val_evaluator
auto_scale_lr = dict(base_batch_size=192 * 8)
从继承的base可以看出,该模型在多个数据集上进行了训练。由于我们只需要在CCPD数据集上进行训练,因此只需继承一个CCPD训练集的配置文件。在param_scheduler
中,有两个学习率衰减的参数,可能是因为该模型在2个数据集上进行了训练,在6个数据集上进行了测试。官方给出的资料也是如此:https://github.com/open-mmlab/mmocr/blob/1.x/configs/textrecog/abinet/README.md
由于我们只需要在一个数据集上进行训练,因此我们只需使用一个策略即可。不过需要注意,那个列表不可以去掉,不能使用:param_scheduler = dict(type='CosineAnnealingLR', by_epoch=True, T_max=10)
,而应该使用 param_scheduler = [dict(type='CosineAnnealingLR', by_epoch=True, T_max=200)]
。
最后需要修改train_list
和test_list
,其他部分不需要修改,修改后的配置文件如下:
_base_ = [
'../_base_/datasets/cdpp.py',
'../_base_/default_runtime.py',
'../_base_/schedules/schedule_adam_base.py',
'_base_abinet_ccpd.py',
]
load_from = 'https://download.openmmlab.com/mmocr/textrecog/abinet/abinet_pretrain-45deac15.pth' # noqa
optim_wrapper = dict(optimizer=dict(lr=1e-4))
train_cfg = dict(max_epochs=200)
# learning policy
# param_scheduler = [
# dict(type='MultiStepLR', milestones=[16, 18], end=20),
# ]
param_scheduler = [dict(type='CosineAnnealingLR', by_epoch=True, T_max=200)]
# dataset settings
train_list = [
_base_.ccpd_textrecog_train]
test_list = [
_base_.ccpd_textrecog_test]
train_dataset = dict(
type='ConcatDataset', datasets=train_list, pipeline=_base_.train_pipeline)
test_dataset = dict(
type='ConcatDataset', datasets=test_list, pipeline=_base_.test_pipeline)
train_dataloader = dict(
batch_size=128,
num_workers=8,
persistent_workers=True,
sampler=dict(type='DefaultSampler', shuffle=True),
dataset=train_dataset)
test_dataloader = dict(
batch_size=64,
num_workers=8,
persistent_workers=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=test_dataset)
val_dataloader = test_dataloader
val_evaluator = dict(
dataset_prefixes=['ccpd'])
test_evaluator = val_evaluator
auto_scale_lr = dict(base_batch_size=192 * 8)
数据集配置文件ccpd.py
和前面差不多,如下:
ccpd_textrecog_data_root = '/root/autodl-tmp/textdrec'
ccpd_textrecog_train = dict(
type='OCRDataset',
data_root=ccpd_textrecog_data_root,
ann_file='textrecog_train.json',
pipeline=None)
ccpd_textrecog_test = dict(
type='OCRDataset',
data_root=ccpd_textrecog_data_root,
ann_file='textrecog_test.json',
test_mode=True,
pipeline=None)
还有一个非常重要的,由于预测的车牌是带有中文的,需要使用特定的字典,否则会报错。我将字典配置在了_base_abinet_ccpd.py
中,所有字典在mmocr的dicts文件夹中,相关内容在 https://mmocr.readthedocs.io/zh_CN/latest/migration/code.html?highlight=字典
官方的文件并没有告诉字典怎么用,我是在mmocr的其他文件中找到用了使用字典的文件,然后抄了一下。
除开字典外,这个文件中还需要配置init_cfg
,在这项里可以加载以前训练过的模型继续训练。手册里给的—resume会报错,只能使用这个方法加载训练过的模型。
配置好的文件如下:
_base_ = '_base_abinet-vision.py'
dictionary = dict( #配置字典
type='Dictionary',
dict_file='{{ fileDirname }}/../../../dicts/chinese_english_digits.txt',
with_padding=True)
model = dict(
decoder=dict(
d_model=512,
num_iters=3,
dictionary=dictionary,
language_decoder=dict(
type='ABILanguageDecoder',
d_model=512,
n_head=8,
d_inner=2048,
n_layers=4,
dropout=0.1,
detach_tokens=True,
use_self_attn=False,
),
init_cfg=dict(
type='Pretrained',
checkpoint='abinet/epoch_20.pth')
),
)
之后训练的过程就和文本检测一样了。
下面是_base_aster_ccpd.py的内容,更改内容已经不太记得了,和原版对比即可。
dictionary = dict(
type='Dictionary',
dict_file='{{ fileDirname }}/../../../dicts/chinese_english_digits.txt',
with_padding=True,
with_unknown=True,
same_start_end=True,
with_start=True,
with_end=True)
model = dict(
type='ASTER',
preprocessor=dict(
type='STN',
in_channels=3,
resized_image_size=(32, 64),
output_image_size=(32, 100),
num_control_points=20),
backbone=dict(
type='ResNet',
in_channels=3,
stem_channels=[32],
block_cfgs=dict(type='BasicBlock', use_conv1x1='True'),
arch_layers=[3, 4, 6, 6, 3],
arch_channels=[32, 64, 128, 256, 512],
strides=[(2, 2), (2, 2), (2, 1), (2, 1), (2, 1)],
init_cfg=[
dict(type='Kaiming', layer='Conv2d'),
dict(type='Constant', val=1, layer='BatchNorm2d'),
]),
encoder=dict(type='ASTEREncoder', in_channels=512),
decoder=dict(
type='ASTERDecoder',
max_seq_len=25,
in_channels=512,
emb_dims=512,
attn_dims=512,
hidden_size=512,
postprocessor=dict(type='AttentionPostprocessor'),
module_loss=dict(
type='CEModuleLoss', flatten=True, ignore_first_char=True),
dictionary=dictionary,
),
data_preprocessor=dict(
type='TextRecogDataPreprocessor',
mean=[127.5, 127.5, 127.5],
std=[127.5, 127.5, 127.5]))
train_pipeline = [
dict(type='LoadImageFromFile', ignore_empty=True, min_size=5),
dict(type='LoadOCRAnnotations', with_text=True),
dict(type='Resize', scale=(256, 64)),
dict(
type='PackTextRecogInputs',
meta_keys=('img_path', 'ori_shape', 'img_shape', 'valid_ratio'))
]
test_pipeline = [
dict(type='LoadImageFromFile'),
dict(type='Resize', scale=(256, 64)),
dict(type='LoadOCRAnnotations', with_text=True),
dict(
type='PackTextRecogInputs',
meta_keys=('img_path', 'ori_shape', 'img_shape', 'valid_ratio',
'instances'))
]
tta_pipeline = [
dict(type='LoadImageFromFile'),
dict(
type='TestTimeAug',
transforms=[[
dict(
type='ConditionApply',
true_transforms=[
dict(
type='ImgAugWrapper',
args=[dict(cls='Rot90', k=0, keep_size=False)])
],
condition="results['img_shape'][1]<results['img_shape'][0]"),
dict(
type='ConditionApply',
true_transforms=[
dict(
type='ImgAugWrapper',
args=[dict(cls='Rot90', k=1, keep_size=False)])
],
condition="results['img_shape'][1]<results['img_shape'][0]"),
dict(
type='ConditionApply',
true_transforms=[
dict(
type='ImgAugWrapper',
args=[dict(cls='Rot90', k=3, keep_size=False)])
],
condition="results['img_shape'][1]<results['img_shape'][0]"),
], [dict(type='Resize', scale=(256, 64))],
[dict(type='LoadOCRAnnotations', with_text=True)],
[
dict(
type='PackTextRecogInputs',
meta_keys=('img_path', 'ori_shape', 'img_shape',
'valid_ratio', 'instances'))
]])
]
下面是aster_resnet45_6e_st_ccpd.py的内容,更改内容已经不太记得了,和原版对比即可。
# training schedule for 1x
_base_ = [
'_base_aster_ccpd.py',
'../_base_/datasets/ccpd.py',
'../_base_/default_runtime.py',
'../_base_/schedules/schedule_adamw_cos_6e.py',
]
# dataset settings
train_list = [
_base_.ccpd_textrecog_train
]
test_list = [
_base_.ccpd_textrecog_test
]
default_hooks = dict(logger=dict(type='LoggerHook', interval=50))
train_dataset = dict(
type='ConcatDataset', datasets=train_list, pipeline=_base_.train_pipeline)
test_dataset = dict(
type='ConcatDataset', datasets=test_list, pipeline=_base_.test_pipeline)
train_dataloader = dict(
batch_size=1024,
num_workers=24,
persistent_workers=True,
pin_memory=True,
sampler=dict(type='DefaultSampler', shuffle=True),
dataset=train_dataset)
auto_scale_lr = dict(base_batch_size=1024)
test_dataloader = dict(
batch_size=1024,
num_workers=4,
persistent_workers=True,
pin_memory=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=test_dataset)
val_dataloader = test_dataloader
val_evaluator = dict(
dataset_prefixes=['ccpd'])
test_evaluator = val_evaluator
测试
实战测试
实战测试即自己加载权重,方法在 https://mmocr.readthedocs.io/zh_CN/dev-1.x/user_guides/inference.html中
网页中的介绍
- 要加载自定义的配置和权重,你可以把配置文件的路径传给
det
,把权重的路径传给det_weights
。>>> MMOCRInferencer(det='path/to/dbnet_config.py', det_weights='path/to/dbnet.pth')
输入可以是以下任意一种格式:
- str: 图像的路径/URL。
>>> inferencer('demo/demo_text_ocr.jpg')
- array: 图像的 numpy 数组。它应该是 BGR 格式。
>>> import mmcv>>> array = mmcv.imread('demo/demo_text_ocr.jpg')>>> inferencer(array)
- **list: 基本类型的列表。列表中的每个元素都将单独处理。`>>> inferencer([‘img_1.jpg’, 'img_2.jpg])
# 列表内混合类型也是允许的>>> inferencer([‘img_1.jpg’, array])`**
- str: 目录的路径。目录中的所有图像都将被处理。
>>> inferencer('tests/data/det_toy_dataset/imgs/test/')
输出是以下字典:
{
'predictions' : [
# 每个实例都对应于一个输入图像
{
'det_polygons': [...], # 2d 列表,长度为 (N,),格式为 [x1, y1, x2, y2, ...]
'det_scores': [...], # 浮点列表,长度为(N, )
'det_bboxes': [...], # 2d 列表,形状为 (N, 4),格式为 [min_x, min_y, max_x, max_y]
'rec_texts': [...], # 字符串列表,长度为(N, )
'rec_scores': [...], # 浮点列表,长度为(N, )
'kie_labels': [...], # 节点标签,长度为 (N, )
'kie_scores': [...], # 节点置信度,长度为 (N, )
'kie_edge_scores': [...], # 边预测置信度, 形状为 (N, N)
'kie_edge_labels': [...] # 边标签, 形状为 (N, N)
},
...
],
'visualization' : [
array(..., dtype=uint8),
]
}
API中展示了更多的参数:
MMOCRInferencer.init():这里面的参数用在初始化中
参数 | 类型 | 默认值 | 描述 |
---|---|---|---|
det | str 或 https://mmocr.readthedocs.io/zh_CN/latest/modelzoo.html#id2, 可选 | None | 预训练的文本检测算法。它是配置文件的路径或者是 metafile 中定义的模型名称。 |
det_weights | str, 可选 | None | det 模型的权重文件的路径。 |
rec | str 或 https://mmocr.readthedocs.io/zh_CN/latest/modelzoo.html#id2, 可选 | None | 预训练的文本识别算法。它是配置文件的路径或者是 metafile 中定义的模型名称。 |
rec_weights | str, 可选 | None | rec 模型的权重文件的路径。 |
kie [1] | str 或 https://mmocr.readthedocs.io/zh_CN/latest/modelzoo.html#id2, 可选 | None | 预训练的关键信息提取算法。它是配置文件的路径或者是 metafile 中定义的模型名称。 |
kie_weights | str, 可选 | None | kie 模型的权重文件的路径。 |
device | str, 可选 | None | 推理使用的设备,接受 torch.device 允许的所有字符串。例如,‘cuda:0’ 或 ‘cpu’。如果为 None,将自动使用可用设备。 默认为 None。 |
[1]: 当同时指定了文本检测和识别模型时,kie
才会生效。
MMOCRInferencer.call():这里的参数用在使用中
参数 | 类型 | 默认值 | 描述 |
---|---|---|---|
inputs | str/list/tuple/np.array | 必需 | 它可以是一个图片/文件夹的路径,一个 numpy 数组,或者是一个包含图片路径或 numpy 数组的列表/元组 |
return_datasamples | bool | False | 是否将结果作为 DataSample 返回。如果为 False,结果将被打包成一个字典。 |
batch_size | int | 1 | 推理的批大小。 |
det_batch_size | int, 可选 | None | 推理的批大小 (文本检测模型)。如果不为 None,则覆盖 batch_size。 |
rec_batch_size | int, 可选 | None | 推理的批大小 (文本识别模型)。如果不为 None,则覆盖 batch_size。 |
kie_batch_size | int, 可选 | None | 推理的批大小 (关键信息提取模型)。如果不为 None,则覆盖 batch_size。 |
return_vis | bool | False | 是否返回可视化结果。 |
print_result | bool | False | 是否将推理结果打印到控制台。 |
show | bool | False | 是否在弹出窗口中显示可视化结果。 |
wait_time | float | 0 | 弹窗展示可视化结果的时间间隔。 |
out_dir | str | results/ | 结果的输出目录。 |
save_vis | bool | False | 是否将可视化结果保存到 out_dir。 |
save_pred | bool | False | 是否将推理结果保存到 out_dir。 |
结语
因为本来只是写给自己的,然后写得可能有些乱,也有些啰嗦(因为自己记性不好,并且还菜)。后来想着在CSDN上白嫖了那么久,应该回馈下自己的一些经验,尽管虽然没那么有用。由于这是我写的第一篇博文,写得不好请见谅!