A Simple and Robust Framework for Cross-Modality Medical Image Segmentation applied to Vision Transformers

TL;DR

This work tackles cross-modality medical image segmentation under data heterogeneity by proposing a simple, registry-free framework that uses Conditional Instance Normalization conditioned on input modality. A single encoder-decoder is trained with interleaved mixed data, and a novel Conditional Vision Transformer (C-ViT) encoder integrates CIN to adapt normalization per modality. The method achieves state-of-the-art cross-modality segmentation on the MM-WHS dataset, outperforming baselines and other cross-modality approaches in both CT target and MRI-assisted settings, with reduced preprocessing and training overhead. The results suggest practical impact for robust, real-world clinical segmentation across diverse imaging modalities and patient datasets.

Abstract

When it comes to clinical images, automatic segmentation has a wide variety of applications and a considerable diversity of input domains, such as different types of Magnetic Resonance Images (MRIs) and Computerized Tomography (CT) scans. This heterogeneity is a challenge for cross-modality algorithms that should equally perform independently of the input image type fed to them. Often, segmentation models are trained using a single modality, preventing generalization to other types of input data without resorting to transfer learning techniques. Furthermore, the multi-modal or cross-modality architectures proposed in the literature frequently require registered images, which are not easy to collect in clinical environments, or need additional processing steps, such as synthetic image generation. In this work, we propose a simple framework to achieve fair image segmentation of multiple modalities using a single conditional model that adapts its normalization layers based on the input type, trained with non-registered interleaved mixed data. We show that our framework outperforms other cross-modality segmentation methods, when applied to the same 3D UNet baseline model, on the Multi-Modality Whole Heart Segmentation Challenge. Furthermore, we define the Conditional Vision Transformer (C-ViT) encoder, based on the proposed cross-modality framework, and we show that it brings significant improvements to the resulting segmentation, up to 6.87\% of Dice accuracy, with respect to its baseline reference. The code to reproduce our experiments and the trained model weights are available at https://github.com/matteo-bastico/MI-Seg.

Figures (6)

• Figure 1: Overview of the proposed cross-modality clinical image segmentation technique applied to the Swin-UNETR segmentation model hatamizadeh_swin_2022. The input modality can be arbitrarily switched to obtain the desired segmentation, while keeping the model unchanged. In conditional models, the modality is also used as input to activate the encoder normalizations corresponding to the data type and generate common latent spaces.
• Figure 2: Overview of the proposed conditional framework: the conditional encoder, $E(\cdot, m)$, generates $S$ common latent spaces, $\{ \mathcal{C}^1, \dots, \mathcal{C}^S\}$, one for each encoder-decoder connection, adapting itself to the input modality $m$ and the decoder, $D(\cdot)$, is unique and does not need adaptation.
• Figure 3: Conditional Vision Transformer Encoder. The multi-head self-attention (MSA) mechanism is followed by a MLP and each of them is preceded by CIN. The latter consists of a switch mechanism to activate the IN corresponding the to modality flag in input. A focus on the different normalizing direction of IN and LN is also shown in the zoom.
• Figure 4: Qualitative comparison of the segmentation result on the CT target domain of our framework applied to the Swin-UNETR model hatamizadeh_swin_2022 on the MM-WHS dataset. Our framework is compared with the baseline model, trained on the target modality, fine-tuning and joint-training with interleaved mixed fashion.
• Figure 5: Comparison of the 3D segmentation of heart substructures from CT using different cross-modality adaptations.