[论文解读] DifFSS: Diffusion Model for Few-Shot Semantic Segmentation

Weimin Tan, Siyuan Chen+1 more|arXiv (Cornell University)|Jul 3, 2023

Domain Adaptation and Few-Shot Learning被引用 13

一句话总结

DifFSS 引入基于扩散模型的范式，用于生成多样化的辅助支持图像以进行少样本语义分割，在不改变现有 FSS 模型架构的前提下提升性能。

ABSTRACT

Diffusion models have demonstrated excellent performance in image generation. Although various few-shot semantic segmentation (FSS) models with different network structures have been proposed, performance improvement has reached a bottleneck. This paper presents the first work to leverage the diffusion model for FSS task, called DifFSS. DifFSS, a novel FSS paradigm, can further improve the performance of the state-of-the-art FSS models by a large margin without modifying their network structure. Specifically, we utilize the powerful generation ability of diffusion models to generate diverse auxiliary support images by using the semantic mask, scribble or soft HED boundary of the support image as control conditions. This generation process simulates the variety within the class of the query image, such as color, texture variation, lighting, $etc$. As a result, FSS models can refer to more diverse support images, yielding more robust representations, thereby achieving a consistent improvement in segmentation performance. Extensive experiments on three publicly available datasets based on existing advanced FSS models demonstrate the effectiveness of the diffusion model for FSS task. Furthermore, we explore in detail the impact of different input settings of the diffusion model on segmentation performance. Hopefully, this completely new paradigm will bring inspiration to the study of FSS task integrated with AI-generated content. Code is available at https://github.com/TrinitialChan/DifFSS

研究动机与目标

• 通过用多样化生成图像丰富支持集来解决少样本分割中单个/少量支持图像的局限性。
• 利用条件扩散模型捕捉类内变异性（颜色、纹理、光照、姿态），从而实现对查询分割的鲁棒性。
• 研究不同的扩散输入条件如何影响分割性能。
• 展示在现有 FSS 架构中集成 DifFSS 的兼容性与性能提升。
• 探讨扩展到 X-shot 的可能性，并讨论生成漂移及对支持质量的敏感性。

提出的方法

• 使用 ControlNet 与 Stable Diffusion，基于支持图像及其分割掩码，通过扩散模型生成辅助支持图像。
• 通过边缘/边界映射和涂鸦（来自 HED 边缘检测）从支持图像创建控制条件。
• 使用诸如 'a real shot photo of {class name}' 的提示来引导图像生成，同时使用与支持图像相同的分割掩码。
• 将生成的辅助图像 I^G 与原始支持图像 I^s 及其分割掩码 M^s 一同输入到标准 FSS 模型 f_seg，以预测查询掩码 M^q_hat。
• 在保持扩散模型参数冻结的同时，使用交叉熵损失训练 f_seg。
• 通过增加辅助样本数量来扩展到 X-shot，并解决 I^G 与 M^s 之间对象位置可能变化所导致的生成漂移。

实验结果

研究问题

• RQ1 将扩散模型以支持分割数据作为条件，是否能产生多样且语义一致的辅助图像，从而提高 FSS 的准确性？
• RQ2 不同的扩散输入条件（分割图、边界图、涂鸦）如何影响分割性能？
• RQ3 DifFSS 方法是否能自然地从 K-shot 扩展到 X-shot，生成漂移的限制有哪些？
• RQ4 生成的辅助样本对 FSS 模型的鲁棒性和类内变异性的表示有何影响？

主要发现

• DiffFSS 在与现有模型结合时，在 PASCAL-5i、FSS-1000 与 MiniCOCO-20i 基准上持续提升了最先进的 FSS 方法。
• 使用扩 diffusion 生成的辅助图像在各数据集上为基线方法（如 BAM、HDMNet）带来显著的 mIoU 提升。
• 通常较多的生成辅助图像会提升性能，但若支持质量差，生成漂移可能抵消收益。
• 控制条件（分割图、边界、涂鸦）均带来性能提升，扩散增强在某些情况下比真实的 5-shot 还获得额外提升。
• 生成图像的原型分布聚集在原图附近，表明生成的图像在扩展类内多样性的同时保持语义一致性。
• 扩展到 X-shot（更多辅助图像）显示出 mIoU 提升，证明 DifFSS 的可扩展性。