《CounTR: Transformer-based Generalised Visual Counting》CVPR2023

摘要

本论文考虑了通用视觉对象计数问题，目标是开发一个计算模型，用于计算任意语义类别的对象数量，使用任意数量的“样本”（即可能为零样本或少样本计数）。作者提出了一个新颖的基于Transformer的架构，称为Counting TRansformer（CounTR），它通过注意力机制显式捕获图像块之间的相似性或与给定“样本”的相似性。此外，作者采用了两阶段训练机制，首先通过自监督学习预训练模型，然后进行有监督的微调。作者还提出了一个简单、可扩展的流程来合成训练图像，这些图像包含大量实例或来自不同语义类别，明确迫使模型利用给定的“样本”。通过在大规模计数基准测试（例如FSC147）上进行彻底的消融研究，证明了在零样本和少样本设置下达到了最先进的性能。

拟解决的问题

论文旨在解决通用视觉对象计数问题，即在一个图像中识别任意语义类别的显著对象的数量。这包括从零样本到少样本对象计数的挑战。

创新之处

1. 提出了一个新颖的基于Transformer的架构CounTR，利用注意力机制显式捕获图像块之间的相似性或与给定样本的相似性。
2. 采用了两阶段训练机制，包括自监督预训练和有监督微调，有效提高了计数性能。
3. 提出了一个简单、可扩展的流程来合成训练图像，这些图像包含大量实例或来自不同语义类别，以解决现有数据集中的长尾问题。
4. 在大规模计数基准测试上进行了彻底的消融研究，证明了模型在零样本和少样本设置下达到了最先进的性能。

方法

• 架构：提出了Counting TRansformer (CounTR)，它包括视觉编码器、特征交互模块（FIM）和视觉解码器。
• 两阶段训练：首先通过掩码图像重建 (MAE) 使用自我监督对模型进行预训练，然后对下游计数任务进行微调。
• 数据合成：提出了一种新颖且可扩展的马赛克管道来合成训练图像，以解决现有对象计数数据集中长尾分布（即具有大量实例的图像往往不太频繁）的挑战。
• 测试时归一化：对于少样本计数，引入了测试时归一化策略来校准输出密度图。

4.1 架构

1. 视觉编码器（Visual Encoder）：将输入图像和用户提供的样本映射到高维特征空间。
2. 特征交互模块（Feature Interaction Module, FIM）：通过一系列的标准Transformer解码器层来融合来自视觉编码器的信息。
3. 视觉解码器（Visual Decoder）：将特征交互模块的输出重新塑形为2D特征图，并恢复到输入图像的原始分辨率。

视觉编码器：由两部分组成：

• ViT编码器：基于Vision Transformer（ViT），将输入图像分割成小块，并通过一系列Transformer编码层处理，输出为一系列D维的特征向量。
• 样本编码器：对于少样本计数，使用轻量级卷积神经网络（CNN）来提取样本的视觉特征。
• 在没有给出示例的零样本场景下，我们采用可学习的 [SPE] 标记作为替代来为模型提供线索。

特征交互模块：FIM由一系列标准的Transformer解码器层构成，其中图像特征作为查询（Query），样本特征或特殊学习到的标记（在零样本情况下）作为键（Key）和值（Value）。这样的设计使得FIM的输出维度与图像特征相同。

视觉解码器：解码器采用渐进式上采样设计，将向量序列首先重塑为密集的特征图，然后通过基于卷积神经网络的解码器进行处理。最终通过一个线性层作为密度回归器，输出单通道的密度热图。

4.2 两阶段训练方案 

1. 自监督预训练：使用遮蔽图像建模（Masked Image Modeling）来训练视觉编码器，通过部分观察到的图像块来重建输入图像。
2. 有监督微调：在预训练的基础上，使用训练集中的图像和样本对模型进行微调，以适应具体的计数任务。

第一阶段：自监督预训练：目的是利用图像内部的冗余信息，通过图像重建任务来学习有效的视觉表示，而不依赖于标注数据。

方法：

• 遮蔽图像建模（Masked Image Modeling, MAE）：这是一种自监督学习方法，通过随机遮蔽图像中的部分区域，然后训练模型预测这些被遮蔽区域的内容。这种方法可以看作是图像版本的BERT，即通过遮蔽部分图像并让模型预测这些部分来学习图像的表示。
• ViT编码器训练：在CounTR模型中，首先使用MAE对基于Vision Transformer（ViT）的视觉编码器进行预训练。具体来说，将输入图像分割成固定大小的块，随机遮蔽一部分图像块，然后使用剩余的图像块作为输入，训练模型来重建整个图像。

第二阶段：在自监督预训练的基础上，通过有监督学习进一步优化模型，使其更好地适应特定的视觉对象计数任务。

方法：

• 模型初始化：使用自监督预训练得到的ViT编码器的权重作为初始化，固定这些权重。
• 任务特定训练：在固定的视觉编码器基础上，添加特征交互模块（FIM）和视觉解码器，然后使用标注的数据集（如FSC-147）对整个模型进行微调。模型的输入包括图像和用户给定的样本，输出为对应的密度图。
• 损失函数：使用像素级均方误差（MSE）作为损失函数，计算模型预测的密度图和真实密度图之间的差异。

 4.3 可扩展的马赛克

解决通用视觉对象计数数据集中的长尾问题，即大多数图像只包含少量对象，而少量图像包含大量对象。方法如下：

• 增加背景多样性：通过结合不同图像的区域，增加训练图像的背景多样性。
• 合成更多实例：在单个图像中合成更多的对象实例，帮助模型更好地学习和泛化到包含大量对象的场景。

拼贴步骤涉及从单个或多个图像中裁剪区域，并将这些区域组合成一个新的图像。 

简单的裁剪和拼贴可能会在图像的边界处留下明显的人工痕迹。为了解决这个问题，在图像的边界处使用随机宽度的α通道混合，使图像的组合更加真实。

4.4 测试时间归一化

针对少样本计数场景提出的一种策略，用于在模型推理阶段校准输出密度图，以提高计数的准确性。 

在少样本计数任务中，用户提供的“样本”（即示例对象）对于模型预测图像中的目标对象至关重要。测试时归一化的目的是利用这些样本信息来校正模型在测试阶段的输出，确保计数的准确性。

1. 样本位置的密度图预测：模型首先对输入图像进行预测，生成一个包含对象位置的密度图。

2. 计算样本位置的平均预测值：在得到的密度图中，计算与用户提供的样本相对应的区域的平均预测值。这些区域通常由样本的边界框（bounding boxes）标识。

3. 归一化处理：如果样本位置的平均预测值超过了预设的阈值（例如1.8），则将整个密度图的值除以该平均值进行归一化。这样做的原因是，如果模型在样本位置的预测值过高，可能意味着它将对象的一部分错误地计为多个实例，或者反之，将多个实例错误地计为一个。

4. 最终计数：经过归一化处理后，通过累加密度图中的所有值来得到最终的对象计数。

在测试时间归一化中，如果密度图中样本位置的平均总和超过 1.8，则密度图的总和将除以该平均值以获得最终的预测。在测试时间裁剪中，如果至少有一个样本的边长小于10个像素，则将图像裁剪成9个片段，模型将分别处理这9幅图像。最终预测将是这 9 张图像的结果之和。