AC-CAM: Affinity-Aware Contrast CAM for Weakly-Supervised Semantic Segmentation on MRI Brain Tumor

Haili Wang, Jingqian Wu, Changwei Wang, Min Zhao, Duzhen Zhang, Rongtao Xu

Abstract: Applying the latest visual transformer (ViT) to Weakly-Supervised Semantic Segmentation (WSSS) can compensate for the local perception limitations of CNN, but it also brings about the over-smoothing problem, that is, the final patch labels tend to be uniform. To overcome this challenge, we present an Affinity-Aware Contrast Class Activation Maps (AC-CAM) framework aimed at enhancing WSSS for MRI Brain Tumor analysis by exploiting only image-level labels. We propose two main components: the Affinity-Aware Token Contrast Module (ATCM) and the Affinity-Aware Refine Module (ARM). ATCM utilizes semantic affinities from attention maps to improve the contrast between patch tokens, effectively reducing the over-smoothing tendency of Vision Transformers (ViT). ARM refines the pseudo labels further, incorporating RGB and affinity information to capture the intricate details of the target objects. Our approach capitalizes on the global feature capturing capabilities of ViT, producing more accurate pseudo-labels for WSSS. The framework is optimized through a composite loss function that ensures the consistency of representations for positive token pairs and discriminability for negative ones. Experiments show that our method achieves state-of-the-art performance on the BraTS 2021 dataset.