《SVTR: Scene Text Recognition with a Single Visual Model》解读

topduke

准程序员

SVTR仅使用单视觉模型，无语言模型，达到与主流视觉-语言两阶段模型相同或者更高的精度，同时，相比与主流的模型，速度更快，参数量更低。由于其优越的性能，SVTR被PP-OCRv3所青睐作为其识别基础模型发布，并且得到了开发者们的广泛关注。这里在SVTR的原始论文的基础上做进一步的解读和伪代码实现，以供学习讨论。

Motivation:

目前高精度的在场景文本识别方法主要有两个模块组成：视觉特征提取模块和序列特征模块用于字符转录，这样两阶段的方法虽然精度高，但是效率低。

效率低主要由语言模型（图1 b、c）的解码器（decoder）引起。有两种方案缓解效率低的问题：

使用并行解码器提高效率，但是会带来一定程度的精度损失；
不使用类语言解码器，仅使用单视觉模型。

SVTR选择了后者。然而，虽然单视觉模型具有高效率的优势，但是精度无法与两阶段方法竞争。SVTR使用单视觉模型，有以下的考量：

以往的单视觉模型仅仅使用图像分类网络如ResNet，MobileNet，ViT等等提取视觉特征，缺乏对场景文本视觉特点的思考，提取的视觉特征是低质量的，无法应对复杂场景的文本识别。

我们认为，单视觉模型需要做到以下几点，才有可能得到场景文本识别的关键特征，应对复杂的文本识别场景，提升识别精度：

1、单视觉模型应具备关注笔画特征的能力

文本字符是由若干个笔画组成，而对于一些相似的字符，如‘B’和‘8’、‘z’和‘2’、‘大’和‘太’等等，它们在形状上相似，但是在笔画上具有细微的差异，因此模型应具备关注字符笔画细节特征的能力，以更加准确的区分相似字符。

2、单视觉模型需捕捉到单个字符内局部特征和字符之间的长距离全局依赖

单个字符存在与整个文本图像中的局部位置，这需要单视觉模型具备提取局部特征的能力，同时在笔画特征的基础上建立单个字符内笔画的相关性；对于文本图像中的所有字符，它们不是相互孤立的，存在内在的联系，因此要求单视觉模型可以建立多个字符的长距离全局依赖；进一步的，图像文本包含了文本区域和非文本区域，因此在建立文本（多个字符）的全局特征基础上，需区分文本区域和非文本区域。

SVTR具体实现

SVTR整体结构遵循了patch-wise image tokenization 的框架，结合文本识别任务的视觉特点，感知笔画的细节特征、字符形态特征和多个字符之间的长距离依赖性。

1. 整体结构

如图2所示， H*W*3 的图像，通过Progressive Overlapping Patch Embedding得到 \frac{H}{4} * \frac{W}{4} * D_{0} 的字符组件（Character Components），这些字符组件用于表示字符笔画。然后，经过高度逐渐下采样的三个阶段，第i阶段由Li个混合块（Mixing Blocks）、合并（Merging）或组合（Combing）操作构成，在不同的尺度上进行特征提取生成一个称为 C 的表示，经过并行线性预测，得到字符序列。

2. Progressive Overlapping Patch Embedding(POPE)

从视觉transformer中的得到灵感，我们将文本图像拆分为若干个patches，每个patches表示了字符的一部分，被称为字符组件（character components），用于表示字符的笔画特征。我们通过两个级联的卷积模块实现Progressive Overlapping Patch Embedding（POPE）如图3（b）， 将输入 H*W*3的图像嵌入到 \frac{H}{4} * \frac{W}{4} * D_{0} 字符组件中 。

视觉transformer中常用的Linear Patch Embedding（图3（a）），使用线性映射，将输入图像嵌入patches中，POPE相比于该方法结合了非线性、交叠特征和渐进式通道扩张的特点，得到更具代表性的字符组件。

3. Mixing Block

前面提到单视觉模型实现高精度文本识别需要两种特征：

局部细粒度特征：笔画特征、局部的字符形态特征、单个字符中笔画的相关性；
全局粗粒度特征：多个字符之间的相关性和文本区域与非文本区域之间的相关性。

因此，我们设计了具有不同感受野的两个混合块（Mixing Block）来完成这两种粒度特征的感知和提取。

Global Mixing：

如图4(a) 所示，Global Mixing感知所有字符组件之间的依赖关系，提取全局特征。此外它还能够削弱非文本区域中组件的影响，进一步提高文本区域中字符组件的重要性。具体地，我们使用 MHSA 和 MLP实现Mixing，整体结构借鉴Vision Transformer Unit。

Local Mixing：

如图4(b) 所示，Local Mixing感知预定义感受野领域（7*11）内组件之间的相关性，提取局部特征。其目标是对字符形态特征进行编码并建立字符内组件之间的关联，从而模拟对字符识别至关重要的笔画特征。

Mixing代码实现：

Local Mixing的灵感来源于卷积，根据卷积的实现方法，Local Mixing的第一种实现方式：

# H、W：特征图的高宽
# kh、kw：局部窗口的高宽
# Wq、Wk、Wv、Wp：可学习参数
# @： 矩阵乘法
# X为[H, W, D]形状的输入，此时X为2D的形状
Q, K, V = X @ (Wq, Wk, Wv)
Q = Reshape(Q, [H*W, 1, D])
KV = Concat(K, V)
KV = Unfold(KV, kernel_sizes=[kh, kw])  
# KV经过Unfold操作后形状为：[2, H*W, kh*kw, D]，Unfold操作提取出所有字符组件kh*kw领域内的元素。
K, V = KV[0], KV[1]
Atten = Softmax(Q @ K) # [H*W, 1, D] @ [H*W, kh*kw, D] = [H*W, 1, kh*kw]
X = Reshape(Atten @ V, [H, W, D]) # [H*W, 1, kh*kw] @ [H*W, kh*kw, D] = [H*W, 1, D] --> [H, W, D]
Output = X @ Wp

然而这种方式由于使用Unfold操作耗时长，我们实现了第二种Global和Local Mixing通用的方式：

# Mixing 实现：
# X为[HW, D]形状的输入，表示有H*W个字符组件，每个字符组件维度为D
Q, K, V = X @ (Wq, Wk, Wv)
if Mixing is Local: 
    Atten = (Q @ K) + Local_mask
elif Mixing is Global:
    Atten = Q @ K
# Atten：[HW, HW]
Atten = Softmax(Atten)
Output = (Atten @ V) @ Wp
--------------------
# Local_mask的生成代码:
Mask = Full([H * W, H + kh - 1, W + kw - 1], -inf)  # 创建元素为无限小的Mask矩阵
h from 0 to H-1:
   w from 0 to W-1:
      Mask[h * W + w, h:h + kh, w:w + kw] = 0. # 将kh*kw领域内的元素置为0
Local_mask = mask[:, kh // 2:H + kh // 2, kw // 2:W + kw //2].flatten(1, 2)