UM-CAM: Uncertainty-weighted multi-resolution class activation maps for weakly-supervised segmentation

Abstract

Weakly-supervised medical image segmentation methods utilizing image-level labels have gained attention for reducing the annotation cost. They typically use Class Activation Maps (CAM) from a classification network but struggle with incomplete activation regions due to low-resolution localization without detailed boundaries. Differently from most of them that only focus on improving the quality of CAMs, we propose a more unified weakly-supervised segmentation framework with image-level supervision. Firstly, an Uncertainty-weighted Multi-resolution Class Activation Map (UM-CAM) is proposed to generate high-quality pixel-level pseudo-labels. Subsequently, a Geodesic distance-based Seed Expansion (GSE) strategy is introduced to rectify ambiguous boundaries in the UM-CAM by leveraging contextual information. To train a final segmentation model from noisy pseudo-labels, we introduce a Random-View Consensus (RVC) training strategy to suppress unreliable pixel/voxels and encourage consistency between random-view predictions. Extensive experiments on 2D fetal brain segmentation and 3D brain tumor segmentation tasks showed that our method significantly outperforms existing weakly-supervised methods. Code is available at: https://github.com/HiLab-git/UM-CAM.