deep-text-recognition-benchmark
将预训练Vision Transformer直接迁移到文字识别任务，发表在ICDAR 2021，支持GPU/CPU/Raspberry Pi多平台推理
加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
微信扫码咨询客服
POMOAI — 让开源 AI 即刻运行
将预训练Vision Transformer直接迁移到文字识别任务，发表在ICDAR 2021，支持GPU/CPU/Raspberry Pi多平台推理
加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
场景文字识别（Scene Text Recognition，STR）是计算机视觉领域的一个硬骨头。当你在街头随手拍一张照片，照片里的路牌、店铺名、广告牌上的文字，形态各异——有水平排列的，有弯曲变形的，有光照不均的，有被遮挡的。传统OCR系统面对这些"野生"场景，往往力不从心。
2021年，一位博士后提出了一个看似"偷懒"的想法：既然Vision Transformer（ViT）已经在图像分类上证明了实力，为什么不直接用它来做文字识别？ 这就是ViTSTR——Vision Transformer for Scene Text Recognition的诞生背景。它基于CLOVA AI的deep-text-recognition-benchmark框架改写，最终发表在ICDAR 2021会议上。
这个项目不是学术"玩具"。它提供了完整的预训练权重，支持从Raspberry Pi到高端GPU的全场景部署，并在GitHub上获得了300+ Stars，说明确实有人真的在用它解决实际问题。
传统STR模型通常采用四阶段流水线：
第一阶段：几何变换（TPS）——用一个薄板样条变换（Thin Plate Spline）将弯曲/倾斜的文字图像"掰正"，类似于PS里的透视矫正。
第二阶段：特征提取——用CNN（VGG/ResNet/RCNN）从矫正后的图像中提取视觉特征，压缩高度维度，把二维图像转成一维特征序列。
第三阶段：序列建模（BiLSTM）——用双向LSTM捕捉字符之间的上下文依赖关系，理解"这个字母后面通常跟什么"。
第四阶段：注意力预测——用注意力机制动态聚焦图像中每个字符的位置，输出最终文本。
ViTSTR则做了一件更激进的事：跳过前三个阶段，直接用一个预训练的ViT模型做端到端识别。 具体来说，它使用了timm库中预训练好的ViT（Vision Transformer）——这些模型原本是在ImageNet上训练的，拥有强大的视觉表征能力。ViTSTR将这些预训练权重迁移到STR任务上，效果出奇地好。
这种"预训练+微调"的思路，本质上和NLP领域的GPT/BERT一脉相承：先让模型在海量数据上学到通用视觉表征，再在特定任务（文字识别）上微调。ViTSTR证明了这一范式在视觉领域同样有效。
ViTSTR提供了三个规格的模型，满足不同场景的性能-速度权衡需求：
| 规格 | 参数量 | 精度 | GPU推理速度 | 树莓派速度 |
|---|---|---|---|---|
| Tiny | 约6M | 82.1% | ~3ms/图 | ~464ms/图 |
| Small | 约22M | 84.2% | ~3ms/图 | ~464ms/图 |
| Base | 约86M | 85.2% | 稍慢 | N/A |
这些数字意味着什么？
在GPU环境下，ViTSTR-Small仅需3毫秒就能识别一张图片，这意味着每秒可以处理300+张图像，已经能够满足视频流实时OCR的需求。而经过INT8量化后，它可以在Raspberry Pi 4上运行，464毫秒一张——这对于嵌入式设备和IoT场景来说已经相当可用。
从基准测试结果来看，ViTSTR的策略非常聪明：虽然精度（85.2%）略低于传统四阶段流水线TRBA（84.3%），但参数量和推理速度大幅优化。在"同等精度下更快"或"同等速度下更准"的赛道上，ViTSTR都有不俗表现。
项目核心代码由5个Python模块构成，分工明确：
model.py —— 统一模型工厂，根据配置参数动态组装四阶段Pipeline或ViTSTR单阶段模型，支持torch.nn.Module的完整序列化（JIT tracing）。
modules/vitstr.py —— ViTSTR核心实现，调用timm.create_vitstr()创建预训练ViT模型，是整个项目的创新核心所在。
modules/feature_extraction.py —— 传统CNN特征提取器（VGG/ResNet/RCNN三选一），当使用ViTSTR时被跳过（--Transformation None --FeatureExtraction None）。
modules/transformation.py —— TPS几何变换网络，实现文字区域矫平。
modules/sequence_modeling.py —— BiLSTM序列建模层。
infer.py —— 推理入口脚本，支持本地模型路径或GitHub URL自动下载预训练权重，还支持JIT编译加速和INT8量化。
train.py —— 训练脚本，支持单卡/多卡分布式训练，需要预先准备好LMDB格式数据集。
依赖项方面，timm库是关键——它提供了预训练Vision Transformer的权重和统一接口；wandb支持训练过程可视化；LMDB作为高效数据库存储大规模训练语料。
对于只是想用ViTSTR做推理的用户来说，体验非常友好。零Docker，纯pip，一行命令出结果：
# 安装
pip3 install -r requirements.txt
# 推理（自动下载预训练权重）
python3 infer.py --image your_image.png \
--model https://github.com/roatienza/deep-text-recognition-benchmark/releases/download/v0.1.0/vitstr_small_patch16_jit.pt
推理效果示例：
→ Available
→ Londen（弯曲文字识别）
→ Greenstead（自然场景文字）
如果想训练自己的模型，则需要额外准备LMDB格式数据集（可从CLOVA AI项目下载），门槛会高一些。
必须承认，ViTSTR不是银弹。它的局限性主要体现在：
1. 不支持任意长度文本。 ViTSTR输出固定长度序列（默认25个token），对于超长文本会截断。这在处理长句子时是个硬伤。
2. 端到端版本有前提。 ViTSTR单阶段Pipeline要求输入图像固定为224×224像素。虽然对主流场景够用，但实际场景中文字图像的长宽比差异很大（短招牌 vs 长横幅），resize到正方形会损失信息。
3. 训练数据依赖CLOVA LMDB。 官方推荐使用CLOVA AI项目预处理好的LMDB数据集，自己准备训练数据有一定门槛。
4. 不支持Transformer编码器+CTC解码的更先进方案。 后续研究（如PARSeq、SWIN-STR）已经超越ViTSTR的架构。如果你需要最新SOTA，可能要另寻他路。
ViTSTR的学术贡献不在于刷榜，而在于证明了轻量化Transformer可以直接迁移到视觉任务，并取得工程上可接受的精度-速度平衡。在它发表之前，社区普遍认为Transformer的计算成本太高、不适合OCR这样的密集预测任务。ViTSTR用实验数据打破了这个认知。
从更宏观的角度看，ViTSTR代表了一个趋势：预训练视觉模型（ViT）正在蚕食传统CNN的地盘。NLP领域的BERT/GPT范式正在向计算机视觉渗透，这是过去几年CV领域最重要的一条暗线。
对于开发者而言，ViTSTR是一个值得收藏的工具：它代码清晰、文档完整、部署友好，在"快速给图片去文字"或"自动化标注场景文字"这类场景下非常实用。如果你的树莓派项目需要文字识别能力，或者你想在视频流里实时OCR，它都是值得一试的选择。