论文笔记 - Towards Weakly-Supervised Text Spotting using a Multi-Task Transformer

Posted on 2022-10-05 Edited on 2022-10-07 In 论文笔记 Views:

1. 前言

为了能够充分利用未标注数据，本文提出了一种基于transformer的端到端文本识别算法，该算法是第一个可以同时使用有监督和弱监督进行训练的端到端文本识别算法，该方法使用同一个隐变量完成多个任务，并提出一种新的适用于文本识别的匈牙利匹配Loss，该方法可以使用只带有文本标注的图片进行训练，并达到与有监督算法相当的指标，加入标注数据后进行，性能会超过现有的有监督算法

论文：Towards Weakly-Supervised Text Spotting using a Multi-Task Transformer

2. 背景

现有的文本识别算法依赖的标注数据需要同时包含文本标注和文本框(box/poly)标注，但是标注成本很高，且文本标注耗费的时间仅为文本框标注时间的一半，因此本文希望可以仅使用文本标注的真实数据进行训练

本文的结构为transformer+多任务头，使用transfomer可以实现通过query获得感兴趣区域的特征，避免依赖区域信息(box/poly)作为输入，此外通过一个object query embedding，经过多个轻量级任务头，可以获得检测、识别、分割结果。

本文贡献如下：

提出一个弱监督训练策略，对真实的标注数据，仅要求文本而不要求空间信息；同时实现了一种适用于文本识别的匈牙利匹配loss
第一个基于transfomer的多任务端到端识别算法，可同时输出检测和识别结果
第一个可以使用同一个架构完成有监督和弱监督训练的算法

3. 本文框架

TTS Framework

本文框架基于Deformable-DETR算法，输入特定数量的query，获得对应数量的embedding，然后输入不同的任务头，获得对应结果，例如：

检测头：由3层的FFN组成，预测文本的bbox
识别头：基于LSTM，预测文本序列
分割头：4个FC+3个反卷积层，预测文本的分割结果(由角点组成)

3.1 匈牙利loss-有监督

训练过程中，网络的预测结果需要和Ground truth进行匹配，本文与Deformable DETR相同，采用匈牙利算法计算一对一的匹配，如下所示

\hat{\sigma}=\underset{\sigma \in \theta_N}{\operatorname{argmin}} \sum_{i=1}^N C\left(y_i, \hat{y}_{\sigma(i)}\right),

其中 $\theta_N$ 为可能得匹配关系数量， $C$ 为cost函数， $y_i$ 为gt， $\hat{y}$ 为网络的预测结果。得到最佳匹配后，计算的loss如下：

L_{\text {Hungarian }}(y, \hat{y})=\sum_{i=1}^N L\left(y_i, \hat{y}_{\hat{\sigma}(i)}\right)

对于文本识别而言，需要在Deformable DETR考虑的物体框和物体种类的匹配关系以外，需要再加上文本的匹配，即

\begin{array}{r} C\left(y, \hat{y}_{\sigma(i)}\right)=-\alpha_c \hat{p}_{\sigma(i)}\left(c_i\right)+\mathbb{1}_{\left\{c_i \neq \emptyset\right\}} \alpha_{\mathrm{box}} C_{\mathrm{box}}\left(b_i, \hat{b}_{\sigma(i)}\right) \\ +\mathbb{1}_{\left\{c_i \neq \emptyset\right\}} \alpha_{\mathrm{rec}} C_{\mathrm{rec}}\left(t_i, \hat{t}_{\sigma(i)}\right) \end{array}

其中， $\hat{p}_{\sigma(i)}\left(c_i\right)$ 为预测结果属于 $c_i$ 类的概率，最终的loss可以表示为：

\begin{aligned} &\quad L\left(y_i, \hat{y}_{\hat{\sigma}(i)}\right)=-\beta_c \log _{\hat{\sigma}(i)}\left(c_i\right) \\ &+\mathbb{1}_{\left\{c_i \neq \emptyset\right\}} \beta_{\mathrm{box}} L_{\mathrm{box}}\left(b_i, \hat{b}_{\hat{\sigma}(i)}\right)+\mathbb{1}_{\left\{c_i \neq \emptyset\right\}} \beta_{\mathrm{rec}} L_{\mathrm{rec}}\left(t_i, \hat{t}_{\hat{\sigma}(i)}\right) . \end{aligned}

在上述公式中，物体框和物体种类的cost与loss与Deformable DETR保持一致，文本识别cost和loss都为cross entropy loss。

如果不考虑文本识别cost，Deformable DETR倾向与选择IoU最大的预测结果，但文本不一定完整，在加入文本识别loss后可以较好的改善文本匹配结果

3.2 匈牙利loss-弱监督

当使用仅包含文本信息的标注信息进行训练时，则需要去除bbox相关的cost及loss，及使用如下公式：

\begin{aligned} C_{\text {weak }}\left(y, \hat{y}_{\sigma(i)}\right)=-\alpha_c \hat{p}_{\sigma(i)}\left(c_i\right)+\mathbb{1}_{\left\{c_i \neq \emptyset\right\}} \alpha_{\mathrm{rec}} C_{\mathrm{rec}}\left(t_i, \hat{t}_{\sigma(i)}\right) \end{aligned}

\begin{aligned} &L_{\text {weak }}\left(y_i, \hat{y}_{\hat{\sigma}(i)}\right)=\quad-\beta_c \log \hat{p}_{\hat{\sigma}(i)}\left(c_i\right)+\mathbb{1}_{\left\{c_i \neq \emptyset\right\}} \beta_{\mathrm{rec}} L_{\mathrm{rec}}\left(t_i, \hat{t}_{\hat{\sigma}(i)}\right) \end{aligned}

4. 实验

本文首先在合成数据上通过有监督的方式训练得到 $TTS_{synthetic}$ ，然后基于该模型分别在全标注和半标注的数据上进行finetune，分别得到 $TTS_{box}$ 和 $TTS_{weak}$ ，为了得到分割结果，通过冻结 $TTS_{box}$ 的权重并使用poly标注数据进行训练，得到 $TTS_{poly}$