【MMOCR】使用 MMOCR 训练自己的文本检测模型

MMOCR 是基于 PyTorch 和 mmdetection 的开源工具箱，专注于文本检测，文本识别以及相应的下游任务，如关键信息提取。

推荐使用 1.x 版本：https://github.com/open-mmlab/mmocr/tree/1.x

一、安装

conda create -n open-mmlab python=3.8 pytorch=1.10 cudatoolkit=11.3 torchvision -c pytorch -y
conda activate open-mmlab
pip3 install openmim
mim install mmengine
mim install 'mmcv>=2.0.0rc1'
mim install 'mmdet>=3.0.0rc0'
git clone https://github.com/open-mmlab/mmocr.git
cd mmocr
git checkout 1.x
pip3 install -e .

二、数据准备

MMOCR 中提供了一站式数据集准备方式，可以使用一行命令实现让多个数据集准备就绪。

文档路径

三、支持的方法

3.1 文本检测

• DBNet (AAAI’2020) / DBNet++ (TPAMI’2022)
• Mask R-CNN (ICCV’2017)
• PANet (ICCV’2019)
• PSENet (CVPR’2019)
• TextSnake (ECCV’2018)
• DRRG (CVPR’2020)
• FCENet (CVPR’2021)

3.2 文本识别

• ABINet (CVPR’2021)
• CRNN (TPAMI’2016)
• MASTER (PR’2021)
• NRTR (ICDAR’2019)
• RobustScanner (ECCV’2020)
• SAR (AAAI’2019)
• SATRN (CVPR’2020 Workshop on Text and Documents in the Deep Learning Era)
• SegOCR (Manuscript’2021)

3.3 关键信息提取

• SDMG-R (ArXiv’2021)

3.4 命名实体识别

• Bert-Softmax (NAACL’2019)

四、训练和测试

1、对单张图进行推理并可视化

python mmocr/utils/ocr.py demo/demo_text_ocr.jpg --print-result --imshow

2、训练

# 用 seg 方法来训练文本识别任务
python tools/train.py configs/textrecog/seg/seg_r31_1by16_fpnocr_toy_dataset.py --work-dir seg

# 训练 icdar2015，需要将数据集和 json 放到 data_root 下
python tools/train.py configs/textdet/dbnet/dbnet_r18_fpnc_1200e_icdar2015.py --work-dir dbnet

3、测试

# 测试 dbnet/latest.pth 的 hmean-iou 指标
python tools/test.py configs/textdet/dbnet/dbnet_r18_fpnc_1200e_icdar2015.py dbnet/latest.pth --eval hmean-iou