PaperlayerView exact revision on arXiv

Source-linked AI summary

MCSeg: Pre-training and Fine-tuning Volumetric Pyramid Transformer for Multi-modal Cardiac Image Segmentation

Zhiyu Ye, Hairong Zheng, Tong Zhang

arXiv:2608.30371v1cs.CV

TL;DR

Cardiac segmentation must generalize across modalities despite scarce finely annotated 3D data. MCSeg combines self-supervised ViT pre-training, an SFP bridge to a CNN decoder, and RMI-based fine-tuning, achieving SOTA results across four datasets and strong few-shot performance.

  • Problem

    Scarce finely annotated 3D data, rare-disease imbalance, and CNN domain-generalization limits hinder robust cardiac segmentation across diverse datasets.

  • Method

    MCSeg pre-trains a volumetric ViT with masked autoencoders, remaps its output through an SFP to a CNN decoder, and uses RMI loss during fine-tuning.

  • Results

    MCSeg achieves SOTA performance across four cardiac datasets and demonstrates cross-modality adaptability and few-shot learning potential.

  • Takeaways & Limitations

    MCSeg provides a cardiac-specific pipeline that transfers anatomical priors across CT and MRI tasks and adapts to scarce annotated data.

  • Takeaways & Limitations

    Large-scale 3D MAE pre-training requires substantial computational and memory resources, including hundreds of GPU hours on advanced hardware.

Abstract

from arXiv · show

Automatic cardiac image segmentation is pivotal for diagnosing and treating cardiac diseases. In this work, we introduce MCSeg, a volumetric transformer-based network tailored for multi-modal cardiac segmentation. To overcome the architectural mismatch inherent in existing hybrid networks, we propose a novel Scaling Feature Pyramid (SFP). Unlike conventional skip connections, the SFP effectively bridges the single-scale 3D Vision Transformer (ViT) encoder and the multi-scale CNN decoder by transforming the ViT's output into a hierarchical feature pyramid, ensuring that global contextual information is effectively leveraged. For the training paradigm, the ViT encoder first undergoes self-supervised pre-training via masked image modeling. Subsequently, the network is fine-tuned on downstream tasks, during which a regional mutual information (RMI) loss is integrated to improve boundary segmentation accuracy. In experiments, MCSeg consistently outperforms eleven SOTA methods on CT dataset ImageCHD, multi-modal dataset MM-WHS, MRI dataset HVSMR-2.0 and MSD Heart, highlighting the effectiveness of our MCSeg for multi-modal cardiac segmentation tasks. Furthermore, MCSeg's superior performance in few-shot experiment showcases its significant potential in adapting to limited data scenarios. Codes and pre-trained ViT-B weights are open-sourced at https://openi.pcl.ac.cn/OpenMedIA/MCSeg

I. INTRODUCTION

MCSeg targets cardiac segmentation across modalities by combining volumetric transformer representations with self-supervised pre-training and specialized fine-tuning. The approach addresses limited annotated 3D data and seeks robust adaptation across diverse cardiac tasks.

  • Cardiac image segmentation supports non-invasive assessment, diagnosis, treatment planning, disease monitoring, and prognosis.
  • Limited finely annotated 3D data, rare-disease imbalance, and CNN domain-generalization limits complicate training universally applicable cardiac segmentation models.
  • Self-supervised masked image modeling and ViTs are used to exploit unlabeled images and model long-range dependencies for diverse 3D cardiac tasks.
  • MCSeg bridges a single-scale ViT encoder and multi-scale decoder with a Scaling Feature Pyramid to exploit pre-trained 3D representations.
  • The two-stage training paradigm combines self-supervised 3D MAE pre-training, downstream fine-tuning, and RMI loss for boundary segmentation accuracy.
  • MCSeg achieves SOTA performance across ImageCHD, HVSMR-2.0, MM-WHS, and MSD, with reported cross-modality adaptability and few-shot learning potential.

II. RELATED WORKS

Prior cardiac segmentation methods include CNNs, pure transformers, hybrid networks, and self-supervised vision-transformer approaches. MCSeg builds on these directions by combining transformer pre-training with a feature-pyramid interface for CNN decoding.

  • U-Net established the symmetric encoder-decoder paradigm, while volumetric transformers extend segmentation through pure-transformer and hybrid CNN-transformer architectures.
  • Hybrid methods connect transformer encoders to CNN decoders through mechanisms such as deformable transformers or multi-resolution skip connections.
  • Masked image modeling methods such as BEiT, SimMIM, and MAE pre-train vision representations by reconstructing masked image patches.
  • Feature Pyramid Networks construct semantically rich multi-scale features through top-down pathways and lateral connections for pixel-level vision tasks.

A. Network Architecture

MCSeg uses a non-hierarchical volumetric ViT encoder, an SFP that remaps its single-scale output into multiple resolutions, and a CNN decoder for segmentation.

  • MCSeg comprises a volumetric transformer encoder, Scaling Feature Pyramid, and CNN decoder.
  • Unlike hierarchical U-shaped encoders, the architecture extracts global contextual representations with a pre-trained ViT and remaps them into a multi-scale pyramid.
  • The ViT receives non-overlapping cubic patches from a 3D volume, projects them into latent embeddings, adds positional embeddings, and processes them with a 12-layer ViT-B configuration.
  • The SFP converts the encoder’s single-scale output at 1/16 resolution into feature maps at 1/4, 1/8, 1/16, and 1/32 scales.
  • The CNN decoder progressively upsamples low-resolution features and concatenates them with corresponding SFP features, while the final block also uses the original input image.

B. Fine-tuning for Cardiac Segmentation

During fine-tuning, the frozen pre-trained ViT encoder feeds the SFP and decoder, which are trained with a compound loss combining Dice, cross-entropy, and regional mutual information. RMI adds local spatial constraints and is smoothly introduced during training.

  • The pre-trained ViT backbone remains frozen, while only the SFP and decoder are trained to produce segmentation results.
  • The compound objective combines soft Dice, cross-entropy, and RMI losses.
  • RMI complements Dice and cross-entropy by enforcing local consistency between predicted and ground-truth regions for smoother boundaries.
  • The number of input regions is determined by the down-sampling factor ds and square region size r.
  • The RMI contribution is regulated by a smooth step function using the current and maximum training epochs.

IV. EXPERIMENTS

The experiments use ImageCAS for pre-training and ImageCHD for downstream congenital-heart-disease segmentation, with dataset-specific annotations and splits.

  • ImageCAS contains 1,000 CTA images and is used entirely for pre-training because complete heart structures are visible.
  • ImageCHD contains 110 3D CT images spanning 16 congenital heart diseases, with seven annotated cardiac substructures.
  • ImageCHD is split into 77 training, 11 validation, and 22 test cases, while training and validation images also support pre-training.

2) ImageCHD:

The broader evaluation covers cardiovascular MRI, multimodal CT/MRI, and left-atrium MRI segmentation datasets, alongside implementation and evaluation details.

  • HVSMR-2.0 provides 60 cardiovascular MRI scans with masks for four chambers and four great vessels.
  • MM-WHS contains 60 CT and 60 MRI images, with 20 annotated cases per modality and remaining unlabeled images used for pre-training.
  • MSD Task02 evaluates left-atrium segmentation using 20 labeled MRI training cases.
  • The model is implemented in PyTorch and MONAI, pre-trained on four 80G A100 GPUs, and fine-tuned on a single 32G V100 GPU.
  • Experimental results are measured with the Dice coefficient.

C. Comparative Experiments

MCSeg is compared with eleven volumetric segmentation networks across full-data and few-shot cardiac experiments. It achieves strong overall performance and benefits substantially from pre-training under limited data.

  • Comparative setup: The comparison includes pure CNN, pure transformer, and hybrid CNN–Transformer networks, with selected models using original or shared pre-trained weights.
  • Full-data evaluation: On ImageCHD full-data evaluation, MCSeg achieves the highest overall score and leads in five of seven substructures.
  • Few-shot evaluation: The few-shot setting randomly selects eight training samples while keeping the test set unchanged.
  • Few-shot evaluation: Pre-training improves relative Dice scores by 43% for UNETR, 46% for CardiacSeg, and 58% for MCSeg compared with training from scratch.
  • Few-shot evaluation: MCSeg reaches Dice scores of 0.896 ± 0.024 in full-data and 0.776 ± 0.065 in few-shot evaluation, with the smallest performance drop.

2) Whole-heart Segmentation on HVSMR-2.0 Dataset:

On HVSMR-2.0, MCSeg achieves strong segmentation performance and shows qualitative advantages over competing models, particularly on MRI images.

  • Qualitative comparison: MCSeg segmentation results closely align with the ground truth, whereas other models show varying missegmentations across different views.The qualitative comparison is reported for an HVSMR-2.0 case.
  • Cross-modality comparison: MCSeg shows a greater advantage over other models on MRI images than on CT images.This comparison is explicitly reported after evaluating the HVSMR-2.0 results.
  • Evaluation design: The evaluation includes five-fold cross-validation on multimodal datasets using models trained from scratch and with pre-trained weights.The referenced cross-validation results compare eight models across various datasets.

5) Computational Efficiency:

The reported experiments examine cross-validation, qualitative results, computational costs, and architectural ablations for MCSeg and competing models.

  • Cross-validation: MCSeg achieves the highest Dice scores and lowest intra-group variance in nearly all MM-WHS cross-validation folds.This result is reported for both CT and MRI evaluations, with models trained from scratch and with pre-trained weights.
  • Computational efficiency: MCSeg requires fewer trainable parameters, uses less GPU memory, and achieves the shortest inference time among compared methods.Measurements were obtained on ImageCHD with an input size of (128, 128, 128).
  • Ablation studies: The ablation experiments assess SFP configuration, pre-training data scale and modality, fine-tuning strategy, and loss-function combinations.These experiments target the main architectural and training choices in MCSeg.
  • SFP ablation: The proposed SFP outperforms a conventional FPN by 1.5%.Additional SFP analysis finds that constructing the pyramid from the deepest ViT block z12 with the selected scaling configuration gives optimal performance.
  • Pre-training ablation: Pre-trained weights improve Dice scores by over 2.3% on ImageCHD and by 7.0% to 9.5% on two MRI datasets compared with training from scratch.The reported comparison also considers pre-training computational costs and identifies ViT-B223 as the optimal trade-off.

2) Data Scale and Modalities for Pre-training:

The ablations compare downstream training strategies and loss combinations while also evaluating pre-training data scales, modalities, and computational costs.

  • Fine-tuning strategy: Frozen-encoder fine-tuning achieves the highest Dice scores on both ImageCHD and HVSMR-2.0.Among the three compared training strategies, it is reported as the most effective for downstream adaptation.
  • Loss functions: The Dice, CE, and RMI combination achieves a Dice score of 0.892 and an HD95 of 13.4.It is compared with alternative loss-function combinations on ImageCHD.
  • Loss functions: RMI loss refines boundary accuracy beyond the standard Dice and CE combination.The reported evidence links the combined loss to improved boundary-sensitive performance.

E. Discussion

MCSeg consistently outperforms existing methods across CT and MRI cardiac segmentation, with transfer benefits from CT pre-training and strong few-shot performance. Its gains combine large-scale cardiac pre-training with SFP and RMI, but deployment remains constrained by computational demands and cardiac-specific scope.

  • Performance: MCSeg outperforms SOTA methods across four cardiac datasets and consistently surpasses existing methods on both CT and MRI datasets.The evaluated datasets include ImageCHD, HVSMR-2.0, MM-WHS, and MSD.
  • Cross-modality transfer: CT-only pre-training substantially improves both CT and MRI performance because the modalities share anatomical structures.The benefit is especially pronounced for MRI tasks affected by scanner and protocol variability.
  • Ablation and architecture: The SFP module’s effectiveness is isolated by applying identical pre-trained ViT weights to UNETR, after which MCSeg remains consistently superior.This comparison directly evaluates the encoder-decoder bridging design rather than changing the pre-training variant.
  • Limitations and accessibility: Large-scale 3D MAE pre-training requires hundreds of GPU hours and advanced hardware, limiting practical accessibility despite released weights and code.The authors open-source pre-trained and fine-tuned weights plus the complete source code to support deployment and reproducibility.
  • Design implications: MCSeg combines pre-trained global anatomical representations with SFP and RMI loss to improve segmentation precision and boundary detail.Large-scale pre-training raises the performance floor, while SFP and RMI loss further improve precision.
  • Data efficiency: Few-shot experiments demonstrate MCSeg’s adaptability when annotated cardiac data are scarce.This supports the model’s use in limited-data segmentation settings.
Loading 2608.30371v1…