MedSAM/MedSAM2 Feature Fusion: Enhancing nnUNet for 2D TOF-MRA Brain Vessel Segmentation

Abstract

Accurate segmentation of brain vessels is critical for diagnosing cerebral stroke, yet existing AI-based methods struggle with challenges such as small vessel segmentation and class imbalance. To address this, our study proposes a novel 2D segmentation method based on the nnUNet framework, enhanced with MedSAM/MedSAM2 features, for arterial vessel segmentation in time-of-flight magnetic resonance angiography (TOF-MRA) brain slices. The approach first constructs a baseline segmentation network using nnUNet, then incorporates MedSAM/MedSAM2's feature extraction module to enhance feature representation. Additionally, focal loss is introduced to address class imbalance. Experimental results on the CAS2023 dataset demonstrate that the MedSAM2-enhanced model achieves a 0.72% relative improvement in Dice coefficient and reduces HD95 (mm) and ASD (mm) from 48.20 mm to 46.30 mm and from 5.33 mm to 4.97 mm, respectively, compared to the baseline nnUNet, showing significant enhancements in boundary localization and segmentation accuracy. This approach addresses the critical challenge of small vessel segmentation in TOF-MRA, with the potential to improve cerebrovascular disease diagnosis in clinical practice.

Keywords: MedSAM; MedSAM2; TOF-MRA; brain vessel segmentation; nnUNet.

Figure 1

Framework overview. MedSAM/MedSAM2 embeddings fuse with the nnUNet encoder‘s specific stages, then progress through the decoder ($D_0$–$D_n$) with skip connections for mask prediction. The green arrow indicates the impending initiation of feature fusion.

Figure 2

Encoder architectures comparison: (a) MedSAM image encoder includes patch processing, embedding, multi-layer transformer blocks (with LayerNorm, multi-head self-attention, and MLP), and feature output. (b) MedSAM2 image encoder uses a Hiera backbone with hierarchical attention (e.g., Q-Pooling and window attention) for multi-scale feature extraction.

Figure 3

Architecture overview of nnUNet. The model processes $640\times 640$ input images through stacked convolutional blocks ($E_0$–$E_7$) for hierarchical feature extraction, with skip connections and transposed convolutions enabling precise segmentation through U-Net encoder–decoder structure. Deep supervision ($L_0$–$L_n$) is applied at multiple levels.

Figure 4

Schematic of nnUNet’s stacked convolutional blocks in encoder–decoder architecture. (Top) Encoder’s downsampling StackedConvBlocks ($E_n$). (Bottom) Decoder’s StackedConvBlocks ($D_n$). Blue dashed frames demarcate modular units, with arrows indicating feature flow directions.

Figure 5

Comparison of two feature fusion methods for MedSAM2 and nnUNet. (Top) MedSAM2 features fused at both $E_3$ and $E_7$ encoder stages of nnUNet. (Bottom) MedSAM features fused only at $E_7$ stage. Both pipelines process input images ($640\times 640$) through independent encoding paths with downsampling, channel adjustment (Channel Up), and feature enhancement (FrequencyLoRA), followed by CAT and projection operations.

Figure 6

Schematic of the FrequencyLoRA module. (Left) FrequencyAdapter processes input (B,C,H,W) through FFT2 spectrum analysis and MLP-based feature enhancement (green dashed frame). (Right) LoRAAdapter performs low-rank adaptation via channel down/up operations (yellow dashed frame). Both modules employ residual connections (dotted arrows) to maintain original features (see Section 2.4 for implementation details).

Figure 7

Architecture of the AttentionGate module. The input tensor $(B,C,H,W)$ undergoes parallel processing: (1) channel attention (blue path) via global average pooling and two-layer MLP and (2) spatial attention (orange path) through feature compression and spatial modeling. Both branches output broadcast-compatible weights that are fused through element-wise multiplication, preserving the original tensor dimensions. Dashed frames distinguish processing stages while arrows indicate data flow.

Figure 8

TOF-MRA slices and their spatial attention heatmaps in the AttentionGate module.

Figure 9

Preprocessing pipeline for MRA images in nnUNet.

Figure 10

The segmentation performance of nnUNet and nnUNet-MedSAM/MedSAM2: correctly segmented pixels are shown in green, false negative pixels are in light purple, and false positive pixels are in yellow.