Contrastive Learning with Diffusion Features for Weakly Supervised Medical Image Segmentation

TL;DR

This paper introduces CLDF, a novel contrastive learning approach that leverages diffusion features and gradient maps from conditional diffusion models to improve weakly supervised medical image segmentation, addressing noise and boundary issues.

Contribution

The paper proposes a new method combining contrastive learning with diffusion features and gradient maps to enhance segmentation accuracy in weakly supervised medical imaging.

Findings

Significant performance improvements over existing baselines.

Effective reduction of false positives and negatives.

Robust segmentation across multiple medical datasets.

Abstract

Weakly supervised semantic segmentation (WSSS) methods using class labels often rely on class activation maps (CAMs) to localize objects. However, traditional CAM-based methods struggle with partial activations and imprecise object boundaries due to optimization discrepancies between classification and segmentation. Recently, the conditional diffusion model (CDM) has been used as an alternative for generating segmentation masks in WSSS, leveraging its strong image generation capabilities tailored to specific class distributions. By modifying or perturbing the condition during diffusion sampling, the related objects can be highlighted in the generated images. Yet, the saliency maps generated by CDMs are prone to noise from background alterations during reverse diffusion. To alleviate the problem, we introduce Contrastive Learning with Diffusion Features (CLDF), a novel method that uses contrastive learning to train a pixel decoder to map the diffusion features from a frozen CDM to a low-dimensional embedding space for segmentation. Specifically, we integrate gradient maps generated from CDM external classifier with CAMs to identify foreground and background pixels with fewer false positives/negatives for contrastive learning, enabling robust pixel embedding learning. Experimental results on four segmentation tasks from two public medical datasets demonstrate that our method significantly outperforms existing baselines.