新闻

• 2024-02-22: 更新支持 PyTorch 2.0 和 Lightning 2.0
• 2024-01-16: 在 NVIDIA 开发者博客中被介绍
• 2023-11-18: 发布了在 ECCV 2022 上与 Deci AI 的采访
• 2023-09-07: 加入了PaddleOCR，这是最受欢迎的多语言 OCR 工具包之一
• 2023-06-15: 加入了docTR，这是一个基于深度学习的 OCR 库
• 2022-07-14: 首次公开发布（在发布时在Papers With Code上 STR 总排名第一）
• 2022-07-04: 被 ECCV 2022 接收

<div align="center">

基于置换自回归序列模型的场景文本识别

Darwin Bautista 和 Rowel Atienza

菲律宾大学迪利曼分校<br/> 电气与电子工程学院

方法 | 示例结果 | 入门指南 | 常见问题 | 训练 | 评估 | 引用

</div>

场景文本识别（STR）模型利用语言上下文来提高对噪声或损坏图像的鲁棒性。最近的方法如 ABINet 使用独立或外部的语言模型（LM）进行预测优化。在这项工作中，我们表明外部 LM（需要预先分配专用计算能力）由于其性能与成本特性较差，对 STR 来说效率不高。我们提出了一种更高效的方法，使用置换自回归序列（PARSeq）模型。查看我们的 ECCV 海报和演示以获得简要概述。

注意： P-S 和 P-Ti 分别是 PARSeq-S 和 PARSeq-Ti 的简称。

方法简介

我们的主要洞见是，通过自回归（AR）模型集成，我们可以统一当前的 STR 解码方法（上下文感知 AR 和上下文无关非 AR）以及双向（完形填空）优化模型：

<div align="center"><img src="https://yellow-cdn.veclightyear.com/835a84d5/a45857b3-ca6b-42b9-8a2a-6d61618a3f0c.png" alt="统一 STR 模型" width="75%"/></div>

单个 Transformer 仅通过改变其注意力掩码就可以实现不同的模型。通过正确的解码器参数化，它可以使用置换语言建模进行训练，从而实现对任意输入上下文子集的任意输出位置进行推理。这种任意解码特性产生了一个_统一的_ STR 模型——PARSeq，它能够进行上下文无关和上下文感知的推理，以及使用双向上下文进行迭代预测优化，无需独立的语言模型。PARSeq 可以被视为具有共享架构和权重的 AR 模型集成： 注意：LayerNorm和Dropout层已省略。[B]、[E]和[P]分别代表序列开始(BOS)、序列结束(EOS)和填充标记。T = 25会产生26个不同的位置标记。位置标记既作为查询向量，也作为输入上下文的位置嵌入。对于[B]，不添加位置嵌入。注意力掩码由给定的排列生成，仅用于上下文-位置注意力。L_ce指的是交叉熵损失。

样本结果

<div align="center">

输入图像	PARSeq-S<sub>A</sub>	ABINet	TRBA	ViTSTR-S	CRNN
<img src="https://yellow-cdn.veclightyear.com/835a84d5/e4a1482b-3bb3-4eb3-81b9-fb04cc510d94.jpg" alt="CHEWBACCA" width="128"/>	CHEWBACCA	CHEWBAGGA	CHEWBACCA	CHEWBACCA	CHEWUACCA
<img src="https://yellow-cdn.veclightyear.com/835a84d5/78d9f372-c485-4499-8d88-dcf1cdb36f80.jpg" alt="Chevron" width="128"/>	Chevrol	Chevro_	Chevro_	Chevr__	Chevr__
<img src="https://yellow-cdn.veclightyear.com/835a84d5/61912595-4291-4afd-a5a2-9be7e8784c2a.jpg" alt="SALMON" height="128"/>	SALMON	SALMON	SALMON	SALMON	SA_MON
<img src="https://yellow-cdn.veclightyear.com/835a84d5/9975eb06-57b6-4af7-8624-a08ef750081d.png" alt="Verbandstoffe" width="128"/>	Verbandsteffe	Verbandsteffe	Verbandstelle	Verbandsteffe	Verbandsleffe
<img src="https://yellow-cdn.veclightyear.com/835a84d5/daef8cb9-6131-4a0b-86d7-85aa794438bf.png" alt="Kappa" width="128"/>	Kappa	Kappa	Kaspa	Kappa	Kaada
<img src="https://yellow-cdn.veclightyear.com/835a84d5/53516bb9-6cf5-4683-a3ad-3cae54bbacaf.jpg" alt="3rdAve" height="128"/>	3rdAve	3=-Ave	3rdAve	3rdAve	Coke

注意：粗体字母和下划线分别表示错误和缺失的字符预测。

</div>

入门

本仓库包含PARSeq的参考实现和复现的模型（统称为场景文本识别模型中心）。版权信息请参见NOTICE。 大部分代码采用Apache License v2.0许可（参见LICENSE），而ABINet和CRNN源代码分别采用BSD和MIT许可（详见相应的LICENSE文件）。

演示

在Hugging Face上提供了一个交互式Gradio演示。该演示使用了此处发布的预训练权重。

安装

要求Python >= 3.9和PyTorch >= 2.0。默认的requirements文件将安装最新版本的依赖项（截至2024年2月22日）。

# 使用特定平台构建。其他PyTorch 2.0选项：cu118, cu121, rocm5.7
platform=cpu
# 为指定的PyTorch平台生成requirements文件
make torch-${platform}
# 安装项目和核心 + 训练 + 测试依赖项。子集：[dev,train,test,bench,tune]
pip install -r requirements/core.${platform}.txt -e .[train,test]

更新依赖版本固定

pip install pip-tools
make clean-reqs reqs  # 重新生成所有requirements文件

数据集

从以下链接下载数据集：

1. LMDB存档，包含MJSynth、SynthText、IIIT5k、SVT、SVTP、IC13、IC15、CUTE80、ArT、RCTW17、ReCTS、LSVT、MLT19、COCO-Text和Uber-Text。
2. LMDB存档，包含TextOCR和OpenVINO。

通过Torch Hub使用预训练模型

可用的模型有：abinet、crnn、trba、vitstr、parseq_tiny、parseq_patch16_224和parseq。

import torch
from PIL import Image
from strhub.data.module import SceneTextDataModule

# 加载模型和图像变换
parseq = torch.hub.load('baudm/parseq', 'parseq', pretrained=True).eval()
img_transform = SceneTextDataModule.get_transform(parseq.hparams.img_size)

img = Image.open('/path/to/image.png').convert('RGB')
# 预处理。模型期望一批形状为(B, C, H, W)的图像
img = img_transform(img).unsqueeze(0)

logits = parseq(img)
logits.shape  # torch.Size([1, 26, 95])，94个字符 + [EOS]符号

# 贪婪解码
pred = logits.softmax(-1)
label, confidence = parseq.tokenizer.decode(pred)
print('解码标签 = {}'.format(label[0]))

常见问题

• 如何训练新语言？请参见问题#5和#9。
• 可以导出为TorchScript或ONNX吗？可以，请参见问题#12。
• 如何在自己的数据集上测试？请参见问题#27。
• 如何微调和/或创建自定义数据集？请参见问题#7。
• val_NED是什么？请参见问题#10。

训练

训练脚本可以训练任何支持的模型。您可以使用命令行覆盖任何配置。有关语法的更多信息，请参阅Hydra文档。使用./train.py --help查看默认配置。

<details><summary>不同训练配置的示例命令</summary><p>

使用预训练权重进行微调

./train.py +experiment=parseq-tiny pretrained=parseq-tiny  # 并非所有实验都有预训练权重

训练模型变体/预配置实验

基本模型配置位于configs/model/中，而变体存储在configs/experiment/中。

./train.py +experiment=parseq-tiny  # 一些示例：abinet-sv, trbc

指定训练的字符集

./train.py charset=94_full  # 其他选项：36_lowercase或62_mixed-case。参见configs/charset/

指定训练数据集

./train.py dataset=real  # 其他选项：synth。参见configs/dataset/

更改常规模型训练参数

./train.py model.img_size=[32, 128] model.max_label_length=25 model.batch_size=384

更改与数据相关的训练参数

./train.py data.root_dir=data data.num_workers=2 data.augment=true

修改 pytorch_lightning.Trainer 参数

./train.py trainer.max_epochs=20 trainer.accelerator=gpu trainer.devices=2

请注意，您可以传递任何Trainer参数，如果该参数在configs/main.yaml中原本没有指定，只需在前面加上+即可。

从检查点恢复训练（实验性功能）

./train.py +experiment=<model_exp> ckpt_path=outputs/<model>/<timestamp>/checkpoints/<checkpoint>.ckpt

评估

测试脚本test.py可用于评估使用本项目训练的任何模型。更多信息，请参见./test.py --help。

PARSeq运行时参数可以使用param:type=value格式传递。例如，可以通过./test.py parseq.ckpt refine_iters:int=2 decode_ar:bool=false调用PARSeq NAR解码。

在基准数据集上进行小写字母数字比较（表6）

./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt  # 或使用发布的权重：./test.py pretrained=parseq

使用不同的评估字符集进行基准测试（表4）

./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt  # 小写字母数字（36字符集）
./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased  # 大小写混合字母数字（62字符集）
./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation  # 大小写混合字母数字 + 标点符号（94字符集）

在更具挑战性的数据集上进行小写字母数字比较（表5）

./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --new

基准模型计算需求（图5）

./bench.py model=parseq model.decode_ar=false model.refine_iters=3

输出标签长度与延迟测量（附录I）

./bench.py model=parseq model.decode_ar=false model.refine_iters=3 +range=true

方向鲁棒性基准测试（附录J）

./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation  # 无旋转
./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation --rotation 90
./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation --rotation 180
./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation --rotation 270

使用训练好的模型从图像中读取文本（附录L）

./read.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --images demo_images/*  # 或使用 ./read.py pretrained=parseq

# 对PARSeq使用NAR解码 + 2次迭代细化
./read.py pretrained=parseq refine_iters:int=2 decode_ar:bool=false --images demo_images/*

调优

我们使用Ray Tune进行学习率的自动参数调优。请参见./tune.py --help。扩展tune.py以支持其他超参数的调优。

./tune.py tune.num_samples=20  # 使用20次试验为PARSeq的默认配置找到最佳LR
./tune.py +experiment=tune_abinet-lm  # 为ABINet的语言模型找到最佳学习率