简介

vedastr 是一个基于 PyTorch 的开源场景文本识别工具箱。它设计灵活，旨在支持场景文本识别任务的快速实现与评估。

特性

• 模块化设计
  我们将场景文本识别框架分解为不同的组件，用户可以通过组合不同的模块轻松构建自定义的场景文本识别框架。

• 灵活性
  vedastr 具备足够的灵活性，可以轻松更换模块内的各个组件。

• 模块扩展性
  很容易将新模块集成到 vedastr 项目中。

• 多框架支持
  该工具箱支持多种流行的场景文本识别框架，例如 CRNN、TPS-ResNet-BiLSTM-Attention、Transformer 等。

• 优异性能
  我们重新实现了 deep-text-recognition-benchmark 中的最佳模型，并获得了更高的平均准确率。此外，我们还实现了一个简单的基线模型（ResNet-FC），其性能也相当不错。

许可证

本项目采用 Apache 2.0 许可证 开源发布。

基准测试与模型库

注意：

• 我们使用 MJSynth(MJ) 和 SynthText(ST) 作为训练数据，并在以下数据集上测试模型：IIIT5K_3000、SVT、IC03_867、IC13_1015、IC15_2077、SVTP、CUTE80。您可以在下方找到这些数据集的下载链接：准备数据。

模型	区分大小写	IIIT5k_3000	SVT	IC03_867	IC13_1015	IC15_2077	SVTP	CUTE80	平均
ResNet-CTC	否	87.97	84.54	90.54	88.28	67.99	72.71	77.08	81.58
ResNet-FC	否	88.80	88.41	92.85	90.34	72.32	79.38	76.74	84.24
TPS-ResNet-BiLSTM-Attention	否	90.93	88.72	93.89	92.12	76.41	80.31	79.51	86.49
Small-SATRN	否	91.97	88.10	94.81	93.50	75.64	83.88	80.90	87.19

TPS：空间变换网络

Small-SATRN：利用二维自注意力机制识别任意形状文本，训练阶段区分大小写，而测试阶段不区分大小写。

平均：所有测试数据集上的平均准确率

区分大小写：如果为真，输出区分大小写并包含常见字符；如果为假，输出不区分大小写，仅包含数字和字母。

安装

需求

• Linux
• Python 3.6+
• PyTorch 1.4.0 或更高版本
• CUDA 9.0 或更高版本

我们已测试过以下操作系统和软件版本：

• 操作系统：Ubuntu 16.04.6 LTS
• CUDA：10.2
• Python 3.6.9
• PyTorch：1.5.1

安装 vedastr

1. 创建并激活 conda 虚拟环境。

conda create -n vedastr python=3.6 -y
conda activate vedastr

2. 按照 官方指南 安装 PyTorch 和 torchvision，例如，

conda install pytorch torchvision -c pytorch

3. 克隆 vedastr 仓库。

git clone https://github.com/Media-Smart/vedastr.git
cd vedastr
vedastr_root=${PWD}

4. 安装依赖项。

pip install -r requirements.txt

准备数据

1. 从 deep-text-recognition-benchmark 下载 LMDB 数据，其中包含训练、验证和评估数据。请注意，我们使用的是由 ASTER 发布的 ST 数据集。

2. 按照如下方式创建 data 目录：

cd ${vedastr_root}
mkdir ${vedastr_root}/data

3. 将下载的 LMDB 数据放入 data 目录中，data 目录的结构应如下所示：

data
└── data_lmdb_release
    ├── evaluation
    ├── training
    │   ├── MJ
    │   │   ├── MJ_test
    │   │   ├── MJ_train
    │   │   └仔MJ_valid
    │   └仔ST
    └仔validation

训练

1. 配置

根据您的需求修改 configs/ 中的配置文件（例如 configs/tps_resnet_bilstm_attn.py）。

2. 运行

# 使用 GPU id 0, 1, 2, 3 进行训练
python tools/train.py configs/tps_resnet_bilstm_attn.py "0, 1, 2, 3" 

默认情况下，快照和日志将生成在 ${vedastr_root}/workdir/name_of_config_file 目录下（您可以在配置文件中指定工作目录）。

测试

1. 配置

根据需要修改配置文件（例如 configs/tps_resnet_bilstm_attn.py）。

2. 运行

# 使用 GPU id 0 和 1 进行测试
./tools/dist_test.sh configs/tps_resnet_bilstm_attn.py path/to/checkpoint.pth "0, 1" 

推理

1. 运行

# 使用 GPU id 0 进行推理
python tools/inference.py configs/tps_resnet_bilstm_attn.py checkpoint_path img_path "0"

部署

1. 按照 官方指南 安装 volksdep。

2. 基准测试（可选）

# 使用 GPU_id 为 0 的 GPU 运行基准测试模型
CUDA_VISIBLE_DEVICES="0" python tools/benchmark.py configs/resnet_ctc.py checkpoint_path out_path --dummy_input_shape "3,32,100"

更多可用参数详见 tools/deploy/benchmark.py。

resnet_ctc 的测试结果如下（测试设备：Jetson AGX Xavier，CUDA：10.2）：

框架	版本	输入形状	数据类型	吞吐量 (FPS)	延迟 (ms)
PyTorch	1.5.0	(1, 1, 32, 100)	fp32	64	15.81
TensorRT	7.1.0.16	(1, 1, 32, 100)	fp32	109	9.66
PyTorch	1.5.0	(1, 1, 32, 100)	fp16	113	10.75
TensorRT	7.1.0.16	(1, 1, 32, 100)	fp16	308	3.55
TensorRT	7.1.0.16	(1, 1, 32, 100)	int8(entropy_2)	449	2.38

3. 将模型导出为 ONNX 格式

# 使用 GPU_id 为 0 的 GPU 将模型导出为 ONNX 格式
CUDA_VISIBLE_DEVICES="0" python tools/torch2onnx.py configs/resnet_ctc.py checkpoint_path --dummy_input_shape "3,32,100" --dynamic_shape

更多可用参数详见 tools/torch2onnx.py。

4. 推理 SDK

详情可参考 FlexInfer。

引用

如果您在研究中使用了本工具箱或基准测试，请引用该项目。

@misc{2020vedastr,
    title  = {vedastr: 场景文本识别工具箱},
    author = {Sun, Jun 和 Cai, Hongxiang 和 Xiong, Yichao},
    url    = {https://github.com/Media-Smart/vedastr},
    year   = {2020}
}

联系方式

本仓库目前由 Sun Jun (@ChaseMonsterAway)、Cai Hongxiang (@hxcai) 和 Xiong Yichao (@mileistone) 维护。

vedastr 快速上手指南

vedastr 是一个基于 PyTorch 的开源场景文本识别工具箱，支持 CRNN、TPS-ResNet-BiLSTM-Attention、Transformer 等多种主流框架，具有模块化设计和高灵活性特点。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 16.04+)
• Python: 3.6 或更高版本
• PyTorch: 1.4.0 或更高版本
• CUDA: 9.0 或更高版本 (如需 GPU 加速)

注意：本项目已在 Ubuntu 16.04.6 LTS, CUDA 10.2, Python 3.6.9, PyTorch 1.5.1 环境下测试通过。国内用户建议使用清华源或阿里源加速 Conda 和 Pip 包的下载。

安装步骤

1. 创建并激活虚拟环境

conda create -n vedastr python=3.6 -y
conda activate vedastr

2. 安装 PyTorch 和 torchvision

建议访问 PyTorch 官网 获取适合您环境的安装命令。以下为通用示例：

conda install pytorch torchvision -c pytorch

(国内用户可尝试使用清华源：conda install pytorch torchvision -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/)

3. 克隆代码仓库

git clone https://github.com/Media-Smart/vedastr.git
cd vedastr
vedastr_root=${PWD}

4. 安装依赖

pip install -r requirements.txt

(国内用户建议指定镜像源：pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple)

5. 准备数据

vedastr 使用 LMDB 格式数据。您需要从 deep-text-recognition-benchmark 下载训练、验证和评估数据（注意：ST 数据集请使用 ASTER 发布的版本）。

下载完成后，在项目根目录下创建 data 文件夹并按以下结构存放数据：

mkdir ${vedastr_root}/data
# 将下载的 lmdb 数据放入 data 目录，最终结构应如下：
# data
# └── data_lmdb_release
#     ├── evaluation
#     ├── training
#     │   ├── MJ
#     │   │   ├── MJ_test
#     │   │   ├── MJ_train
#     │   │   └── MJ_valid
#     │   └── ST
#     └── validation

基本使用

训练模型

修改 configs/ 目录下的配置文件（例如 configs/tps_resnet_bilstm_attn.py）以适配您的需求，然后运行以下命令启动训练：

# 使用 GPU 0, 1, 2, 3 进行训练
python tools/train.py configs/tps_resnet_bilstm_attn.py "0, 1, 2, 3" 

训练产生的快照和日志默认保存在 ${vedastr_root}/workdir/配置文件名 目录下。

测试模型

修改配置文件后，使用以下命令对模型进行评估：

# 使用 GPU 0, 1 进行测试
./tools/dist_test.sh configs/tps_resnet_bilstm_attn.py path/to/checkpoint.pth "0, 1" 

请将 path/to/checkpoint.pth 替换为您实际的模型权重路径。

单张图片推理

使用预训练模型对单张图片进行文本识别：

# 使用 GPU 0 进行推理
python tools/inference.py configs/tps_resnet_bilstm_attn.py checkpoint_path img_path "0"

• checkpoint_path: 模型权重文件路径
• img_path: 待识别的图片路径