BetterNucleiSeg: Better Diffusion Model for Robust Cross-Domain Nuclei Segmentation
摘要
Nuclei segmentation is crucial for quantitative pathology applications, including accurate cell identification and counting. Recent advances in deep learning, from CNNs to Transformers, have improved segmentation accuracy. However, these discriminative methods inherently rely on the distribution of their training data, limiting their generalization in real-world scenarios characterized by variations in staining intensity, magnification levels, and tissue preparation protocols. More recently, diffusion-based vision foundation models have demonstrated strong generalization and have been successfully adapted to various vision tasks. Motivated by this, we propose BetterNucleiSeg, a framework that synergizes Stable Diffusion for image prior encoding and Segment Anything Model (SAM) for uncertainty-aware segmentation refinement. Specifically, BetterNucleiSeg adopts Stable Diffusion as the backbone, fine-tuning only the U-Net and freeze VAE component to retain pretrained image priors while disabling other modules. Additionally, to handle boundary ambiguities and partial occlusions, we introduce a SAM-guided attention mechanism. This mechanism transforms SAM outputs into latent-space masks, dynamically reweighting feature activations and emphasizing error-prone regions. Experiments show that BetterNucleiSeg achieves state-of-the-art performance on both MoNuSeg-2018 and TNBC, and also delivers favorable zero-shot generalization capability on unseen datasets.