MSA-UNet3+: Multi-Scale Attention UNet3+ with New Supervised Prototypical Contrastive Loss for Coronary DSA Image Segmentation
About
Accurate segmentation of coronary Digital Subtraction Angiography (DSA) images is essential for diagnosing and treating coronary artery disease (CAD). Despite advances in deep learning, challenges such as high intra-class variance and class imbalance limit precise vessel delineation. Existing approaches for coronary DSA segmentation cannot effectively address these issues. Furthermore, existing segmentation network encoders do not directly generate semantic embeddings, which could enable the decoder to reconstruct segmentation masks more effectively. We propose a Supervised Prototypical Contrastive Loss (SPCL) that combines supervised and prototypical contrastive learning to enhance coronary DSA image segmentation. The supervised contrastive loss enforces semantic embeddings in the encoder, improving feature differentiation. The prototypical contrastive loss enables the model to focus on the foreground class while alleviating high intra-class variance and class imbalance by concentrating only on hard-to-classify background samples. We implement the proposed SPCL within MSA-UNet3+, a Multi-Scale Attention-Enhanced UNet3+ architecture. The architecture integrates a Multi-Scale Attention Encoder (M-encoder), a Multi-Scale Dilated Bottleneck (MSD-Bottleneck) for multi-scale feature extraction, and a Contextual Attention Fusion Module (CAFM) to preserve fine-grained details while improving contextual understanding. Experiments on a private coronary DSA dataset demonstrate that MSA-UNet3+ outperforms state-of-the-art methods, achieving the highest Dice coefficient and F1-score while significantly reducing ASD and ACD. The framework provides precise vessel segmentation for accurate identification of coronary stenosis and supports informed diagnostic and therapeutic decisions. The code will be released at https://github.com/rayanmerghani/MSA-UNet3plus.