mmocr官方代码 : https://github.com/open-mmlab/mmocr
首先感谢一波大佬们的无私,把代码都开源了,还有辛勤的商汤的大佬们一直维护更新。这次写点自己的使用记录,给同我一样刚入门的萌新们。
常规的配置参考官网的教程Getting Started(再次感谢大佬们搞的这么仔细的教程),然后如果是30系列的显卡,因为只支持cuda11,所以环境配置有点麻烦,我这边是3070,就把自己的配置过程放上来:
# mmocr for 3070
conda create -n open-mmlab python=3.7 -y
conda activate open-mmlab
# install latest pytorch prebuilt with the default prebuilt CUDA version (usually the latest)
conda install pytorch==1.8.0 torchvision torchaudio cudatoolkit=11.1 -c pytorch -c nvidia
# install the latest mmcv-full
pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.8.0/index.html
# install mmdetection
pip install mmdet
# install mmocr
git clone https://github.com/open-mmlab/mmocr.git
cd mmocr
pip install -r requirements.txt
pip install -v -e . # or "python setup.py build_ext --inplace"
export PYTHONPATH=$(pwd):$PYTHONPATH
此后运行代码,可能还会有个cocoapi的小问题,报错 AttributeError: COCO object has no attribute get_cat_ids ,有多种解决方案
# 方法1
git clone https://github.com/open-mmlab/cocoapi.git
cd cocoapi/pycocotools
pip install .
# 方法2
pip uninstall pycocotools
pip install mmpycocotools
然后我们就可以跑个官方demo看看环境是否OK了
python demo/ocr_image_demo.py demo/demo_text_det.jpg demo/output.jpg
官方教程Datasets Preparation介绍得很详细了,我再补充点小细节
我是转换成coco类格式的数据,关于coco数据的格式可以参考Gemfield 的 COCO数据集的标注格式。如果你是标准学术的数据集,官方代码 tools里面包含各种数据互相转换的脚本。数据目录设置都在configs对应的.py文件中设置,如下例子
# 指定数据类型,此类型使.txt的标签文件, 'IcdarDataset' 类使用.json标签文件
dataset_type = 'TextDetDataset'
# 图片目录前缀
img_prefix = 'tests/data/toy_dataset/imgs'
# annotation文件
test_anno_file = 'tests/data/toy_dataset/instances_test.txt'
我使用的是 ‘IcdarDataset’ 类数据,官方config基本也都是此类数据,最重要的是.json格式的annotation文件,可以参考官方给的文件样例,字典格式,keys主要是 “images” ,“categories” 和 “annotations” :
"images"
“images” 的value是一个字典列表,列表每个元素是一个字典,包含图片的信息,例子
{"file_name": "training/0336.png", "height": 1200, "width": 1600, "segm_file": "training/0336.xml", "id": 0}
"categories"
"categories"的value是一个字典列表,即标签类别,因为我们ocr里面都是文本类别,所以只要一类就行,无脑复制粘贴即可
"categories": [{"id": 1, "name": "text"}]
"annotations"
“annotations” 的value也是一个字典列表,列表每个元素是一个字典,即最终读取的ground truth标签,例子
{"iscrowd": 0, "category_id": 1, "bbox": [213, 16, 370, 1163], "area": 168314.0, "segmentation": [[485, 1179, 306, 991, 252, 800, 213, 608, 215, 413, 274, 214, 402, 16, 535, 130, 471, 291, 296, 460, 301, 620, 365, 777, 490, 931, 583, 1089]], "image_id": 0, "id": 0}
根据自己不同的数据,按上面格式写好保存成.json文件,然后configs里面填好对应目录即可。跑代码时,如果能跑,但是没有出现loss,只会保存权重文件,那就是你的数据格式有问题,或者对应的目录设置错了。
此类数据比较简单,annotation每行只要指定 文件名和对应的文字label即可,前提是你要把一个个segmentation的区域抠出来做成单独的图片数据,例子:
train_words/1001724.jpg Chiquita
第一部分是文件的路径,绝对路径相对路径都行,和config中的设置相对应就好了;第二部分是文字的真实label
然后在config文件中定义数据和detection差不多,此处直接’OCRDataset’类型就可以,然后train_prefix 指定图片目录前缀,train_ann_file指定对应的annotation文件位置,test设置也是一样。
dataset_type = 'OCRDataset'
train_prefix = 'data/chinese/'
train_ann_file = train_prefix + 'labels/train.txt'
另外需要注意的就是,文字识别需要一个词汇字典,由dict_file定义,如果是做中文文字识别,官方的sar模型已经有预训练的权重了,可以直接下载然后自己fine-tune一下,效果很好的,mmocr官方简直是直接喂饭,大家只要张嘴就好了…
dict_file = 'data/chineseocr/labels/dict_printed_chinese_english_digits.txt'
其实数据准备好了,其他的都很简单,水到渠成,官方的教程也有十分详细的资料,我就不再赘述了。提下自己的操作,因为官方教程训练和测试都是通过调用 .sh脚本执行,有时候不是很方便,其实看下脚本代码就知道,就是调用tools下的 train.py 等文件,那我们可以直接执行.py文件干活,举个简单的例子:
# 导入官方的模型训练,第一个参数 config文件, --work-dir 指定log 权重保存位置等;--load-from 先load模型再训练;--resume-from 继续训练;--gpus 使用的gpu个数; --gpu-ids 指定使用的gpu id
python ./tools/train.py configs/textrecog/sar/sar_r31_parallel_decoder_chinese.py --work-dir ./results/sar/ --load-from checkpoints/sar_chineseocr.pth --gpus 1 --gpu-ids 4
直接看对应的.py文件,看它需要哪些args就知道各个参数了。一般在训练的时候会有evaluation,在config最后一行进行设置。同时在最开始行,会有_base参数,第一个是定义优化器学习率等,第二个是权重保存间隔等。
evaluation = dict(interval=10, metric='hmean-iou')
_base_ = [
'../../_base_/schedules/schedule_1200e.py',
'../../_base_/default_runtime.py'
]
若有后续再补充…