Skip to content

Latest commit

 

History

History
227 lines (161 loc) · 10.8 KB

README_CN.md

File metadata and controls

227 lines (161 loc) · 10.8 KB

English | 中文

PSENet

PSENet: Shape Robust Text Detection With Progressive Scale Expansion Network

1. 概述

PSENet

PSENet是一种基于语义分割的文本检测算法。它可以精确定位具有任意形状的文本实例，而大多数基于anchor类的算法不能用来检测任意形状的文本实例。此外，两个彼此靠近的文本可能会导致模型做出错误的预测。因此，为了解决上述问题，PSENet还提出了一种渐进式尺度扩展算法（Progressive Scale Expansion Algorithm, PSE）,利用该算法可以成功识别相邻的文本实例[1]。

Figure 1. Overall PSENet architecture

图 1. PSENet整体架构图

PSENet的整体架构图如图1所示，包含以下阶段:

  1. 使用Resnet作为骨干网络，从2，3，4，5阶段进行不同层级的特征提取；
  2. 将提取到的特征放入FPN网络中，提取不同尺度的特征并拼接；
  3. 将第2阶段的特征采用PSE算法生成最后的分割结果，并生成文本边界框。

2. 实验结果

ICDAR2015

模型 环境配置 骨干网络 预训练数据集 Recall Precision F-score 训练时间 ms/step 吞吐量 配置文件 模型权重下载
PSENet D910x8-MS2.0-G ResNet-152 ImageNet 79.39% 84.91% 82.06% 11.544 s/epoch 769.6 83.16 img/s yaml ckpt | mindir
PSENet D910x8-MS2.0-G ResNet-50 ImageNet 76.75% 86.58% 81.37% 4.562 s/epoch 304.138 210.43 img/s yaml ckpt | mindir
PSENet D910x8-MS2.0-G MobileNetV3 ImageNet 73.52% 67.84% 70.56% 2.604 s/epoch 173.604 368.66 img/s yaml ckpt | mindir

SCUT-CTW1500

模型 环境配置 骨干网络 预训练数据集 Recall Precision F-score 训练时间 ms/step 吞吐量 配置文件 模型权重下载
PSENet D910x8-MS2.0-G ResNet-152 ImageNet 73.69% 74.38% 74.04% 67 s/epoch 4466.67 14.33 img/s yaml ckpt | mindir

注释：

  • 环境配置：训练的环境配置表示为 {处理器}x{处理器数量}-{MS模式}，其中 Mindspore 模式可以是 G-graph 模式或 F-pynative 模式。
  • PSENet的训练时长受数据处理部分超参和不同运行环境的影响非常大。
  • 在ICDAR15数据集上，以ResNet-152为backbone的MindIR导出时的输入Shape为(1,3,1472,2624) ，以ResNet-50或MobileNetV3为backbone的MindIR导出时的输入Shape为(1,3,736,1312)。
  • 在SCUT-CTW1500数据集上，MindIR导出时的输入Shape为(1,3,1024,1024) 。

3. 快速上手

3.1 安装

请参考MindOCR套件的安装指南 。

3.2 数据准备

3.2.1 ICDAR2015 数据集

请从该网址下载ICDAR2015数据集，然后参考数据转换对数据集标注进行格式转换。

完成数据准备工作后，数据的目录结构应该如下所示：

.
├── test
│   ├── images
│   │   ├── img_1.jpg
│   │   ├── img_2.jpg
│   │   └── ...
│   └── test_det_gt.txt
└── train
    ├── images
    │   ├── img_1.jpg
    │   ├── img_2.jpg
    │   └── ....jpg
    └── train_det_gt.txt

3.2.2 SCUT-CTW1500 数据集

请从该网址下载SCUT-CTW1500数据集，然后参考数据转换对数据集标注进行格式转换。

完成数据准备工作后，数据的目录结构应该如下所示：

ctw1500
 ├── test_images
 │   ├── 1001.jpg
 │   ├── 1002.jpg
 │   ├── ...
 ├── train_images
 │   ├── 0001.jpg
 │   ├── 0002.jpg
 │   ├── ...
 ├── test_det_gt.txt
 ├── train_det_gt.txt

3.3 配置说明

在配置文件configs/det/psenet/pse_r152_icdar15.yaml中更新如下文件路径。其中dataset_root会分别和data_dir以及label_file拼接构成完整的数据集目录和标签文件路径。

...
train:
  ckpt_save_dir: './tmp_det'
  dataset_sink_mode: False
  dataset:
    type: DetDataset
    dataset_root: dir/to/dataset          <--- 更新
    data_dir: train/images                <--- 更新
    label_file: train/train_det_gt.txt    <--- 更新
...
eval:
  dataset_sink_mode: False
  dataset:
    type: DetDataset
    dataset_root: dir/to/dataset          <--- 更新
    data_dir: test/images                 <--- 更新
    label_file: test/test_det_gt.txt      <--- 更新
...

【可选】可以根据CPU核的数量设置num_workers参数的值。

PSENet由3个部分组成：backbone、neck和head。具体来说:

model:
  type: det
  transform: null
  backbone:
    name: det_resnet152
    pretrained: True    # 是否使用ImageNet数据集上的预训练权重
  neck:
    name: PSEFPN         # PSENet的特征金字塔网络
    out_channels: 128
  head:
    name: PSEHead
    hidden_size: 256
    out_channels: 7     # kernels数量

3.4 训练

  • 后处理

训练前，请确保在/mindocr/postprocess/pse目录下按照以下方式编译后处理代码：

python3 setup.py build_ext --inplace
  • 单卡训练

请确保yaml文件中的distribute参数为False。

# train psenet on ic15 dataset
python tools/train.py --config configs/det/psenet/pse_r152_icdar15.yaml
  • 分布式训练

请确保yaml文件中的distribute参数为True。

# n is the number of GPUs/NPUs
mpirun --allow-run-as-root -n 8 python tools/train.py --config configs/det/psenet/pse_r152_icdar15.yaml

训练结果（包括checkpoint、每个epoch的性能和曲线图）将被保存在yaml配置文件的ckpt_save_dir参数配置的路径下，默认为./tmp_det。

3.5 评估

评估环节，在yaml配置文件中将ckpt_load_path参数配置为checkpoint文件的路径，设置distribute为False，然后运行：

python tools/eval.py --config configs/det/psenet/pse_r152_icdar15.yaml

3.6 MindSpore Lite 推理

请参考MindOCR 推理教程，基于MindSpore Lite在Ascend 310上进行模型的推理，包括以下步骤：

  • 模型导出

请先下载已导出的MindIR文件，或者参考模型导出教程，使用以下命令将训练完成的ckpt导出为MindIR文件:

python tools/export.py --model_name_or_config psenet_resnet152 --data_shape 1472 2624 --local_ckpt_path /path/to/local_ckpt.ckpt
# or
python tools/export.py --model_name_or_config configs/det/psenet/pse_r152_icdar15.yaml --data_shape 1472 2624 --local_ckpt_path /path/to/local_ckpt.ckpt

其中，data_shape是导出MindIR时的模型输入Shape的height和width，下载链接中MindIR对应的shape值见注释。

  • 环境搭建

请参考环境安装教程，配置MindSpore Lite推理运行环境。

  • 模型转换

请参考模型转换教程，使用converter_lite工具对MindIR模型进行离线转换。

  • 执行推理

在进行推理前，请确保PSENet的后处理部分已编译，参考训练的后处理部分。

假设在模型转换后得到output.mindir文件，在deploy/py_infer目录下使用以下命令进行推理：

python infer.py \
    --input_images_dir=/your_path_to/test_images \
    --det_model_path=your_path_to/output.mindir \
    --det_model_name_or_config=../../configs/det/psenet/pse_r152_icdar15.yaml \
    --res_save_dir=results_dir

参考文献

[1] Wang, Wenhai, et al. "Shape robust text detection with progressive scale expansion network." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019.