CounTR:基于Transformer的通用视觉计数.
为了实现通用视觉计数,即给定若干目标示例exemplar时在查询图像query中计数该目标类别的数量;本文提出了一种基于Transformer的模型CounTR,首先通过掩码图像重建(MAE)进行自监督预训练,然后对下游计数任务进行微调。作者进一步提出了一种可扩展的Mosaic数据增强来合成训练图像,以解决现目标计数数据集中长尾分布的挑战。
1. CounTR
CounTR模型的结构如图所示。输入查询图像$X_i$和若干目标示例$S_i^k$,首先通过视觉编码器$\Phi_{VIT-ENC},\Phi_{CNN-ENC}$对两者进行编码,再通过特征交互模块$\Phi_{FIM}$进行信息交互,最后通过视觉解码器$\Phi_{DECODER}$生成密度图$y_i$。
\[y_i = \Phi_{DECODER}(\Phi_{FIM}(\Phi_{VIT-ENC}(X_i),\Phi_{CNN-ENC}(S_i^k)))\]视觉编码器$\Phi_{VIT-ENC}$采用ViT网络,把查询图像$X_i$拆分成$16\times 16$的patch,然后采用12层Transformer编码层。视觉编码器$\Phi_{CNN-ENC}$是一个轻量级的卷积网络架构,包括四个卷积层和一个全局平均池化层,将目标示例$S_i^k$编码为向量。
特征交互模块$\Phi_{FIM}$是由一系列Transformer解码层构成的。图像特征直接作为查询向量,示例特征(或可学习的特殊token)通过两个不同的线性投影映射为值向量和键向量。
视觉解码器$\Phi_{DECODER}$采用渐进式上采样设计。向量序列首先被重塑为特征映射,然后使用了4个上采样块,每个上采样块由一个卷积层和一个2×双线性插值组成。最后一次上采样后采用线性层作为密度回归量,输出一个单通道密度热图。
2. 两阶段训练策略
CounTR模型采用两阶段的训练策略。
- 用MAE进行自监督预训练。首先将图像划分为规则的不重叠的patch,并且采样patch的子集作为ViT编码器的输入。计算出的特征进一步通过一个轻量级的解码器重建输入图像,其中可学习的掩码Token和位置编码的组合用作Query。训练损失定义为重建图像与输入图像在像素空间上的均方误差(MSE)。
- 监督微调。使用预训练的ViT权重初始化图像编码器,并在目标计数任务上微调模型。输出图像对应的密度图中目标的统计数量可以通过对离散密度映射求和得到。使用均方误差来评估预测密度图和真实密度图之间的差异。
3. 可扩展的Mosaic增强
CounTR模型采用可扩展的Mosaic数据增强来合成训练图像,以解决现有计数数据集中的长尾问题。Mosaic数据增强包括拼贴(collage)和混合(blending)两个步骤。
拼贴(collage)首先从图像中裁剪一个随机大小的正方形区域,并将其缩放到统一的大小(如原始图像大小的四分之一)。在重复区域裁剪多次后,将裁剪的区域拼贴在一起,并更新相应的密度图。拼贴有两种不同的形式:只使用一个图像或四个不同的图像。如果只使用一张图像,可以增加图像中包含的目标数量,这对解决长尾问题有很大帮助;如果使用四张不同的图像,可以显著提高训练图像的背景多样性,增强模型区分不同类别物体的能力。在实验中,如果图像中包含的目标数量超过阈值,则使用同一图像进行拼贴;否则使用四个不同的图像进行拼贴。
简单的裁剪和拼贴并不能合成完美的图像,因为在边界之间仍然存在明显的伪影。混合(blending)用比原始图像大小的四分之一稍大的尺寸裁剪图像,用以在边界留下特定的空间用于$α$通道混合。使用随机的$α-通道边界宽度,这使得图像的构图更加逼真。
4. 测试时归一化
CounTR模型引入了测试时间归一化策略来校准输出密度图。引入先验知识,即样本边界框内的物体计数应该恰好为$1.0$,因此任何预测偏差都可以通过密度图除以当前样本边界框的预测计数来校准。
此外,对于具有极小物体(样本边长小于10像素)的图像,采用滑动窗口预测,将图像平均分成9块,并将其缩放到原始尺寸。物体的总数是9个图像的单个计数结果的总和。
5. 实验分析
在FSC-147数据集上CounTR在零样本和少样本计数方面都显著优于以前的方法。
下图展示了计数设置的定性结果。CounTR模型可以很容易地计算出物体的数量并确定它们的位置。最后一幅图像由于边界框的模糊性,错误地选择了眼镜镜片而不是太阳镜进行计数,这可以通过测试时间归一化来纠正。