Distillation Learning Guided by Image Reconstruction for One-Shot Medical Image Segmentation

TL;DR

This work tackles the challenge of one-shot medical image segmentation by introducing a reconstruction-guided distillation framework. It combines a registration-based data augmentation network with a teacher–student distillation paradigm where a reconstruction-focused teacher guides a segmentation-focused student via feature alignment and a cosine-similarity-based hint loss. The approach leverages real unlabeled images to learn anatomy-rich representations and uses a lightweight inference model for efficient deployment, achieving superior generalization across brain MRI, abdominal CT, and vertebrae CT datasets compared to state-of-the-art one-shot methods. The results indicate strong potential for robust MIS in data-scarce and cross-modality scenarios, with implications for broader clinical applicability and downstream diagnostic tasks.

Abstract

Traditional one-shot medical image segmentation (MIS) methods use registration networks to propagate labels from a reference atlas or rely on comprehensive sampling strategies to generate synthetic labeled data for training. However, these methods often struggle with registration errors and low-quality synthetic images, leading to poor performance and generalization. To overcome this, we introduce a novel one-shot MIS framework based on knowledge distillation, which allows the network to directly 'see' real images through a distillation process guided by image reconstruction. It focuses on anatomical structures in a single labeled image and a few unlabeled ones. A registration-based data augmentation network creates realistic, labeled samples, while a feature distillation module helps the student network learn segmentation from these samples, guided by the teacher network. During inference, the streamlined student network accurately segments new images. Evaluations on three public datasets (OASIS for T1 brain MRI, BCV for abdomen CT, and VerSe for vertebrae CT) show superior segmentation performance and generalization across different medical image datasets and modalities compared to leading methods. Our code is available at https://github.com/NoviceFodder/OS-MedSeg.

Figures (7)

• Figure 1: Overview of our problem. Our proposed method achieves natural, realistic, and smooth segmentation, outperforming current state-of-the-art one-shot methods (CLMorph liu2023contrastive, DataAug zhao2019data, and BRBS he2022learning).
• Figure 2: Schematic of our one-shot medical image segmentation framework, consisting of three stages: (a) Generating Labeled Data using image registration and contrastive learning, (b) Feature Distillation Learning through image reconstruction and segmentation, and (c) Inference with a lightweight student network predicting segmentation labels on unknown images. $G$ denotes the registration network, $T$ and $S$ are the teacher and student networks, and $H$ represents the extracted features.
• Figure 3: The proposed teacher-student network uses a residual join U-Net, with each residual block having convolution, PReLU, and layer normalization. It has two output headers: Seg Head for segmentation (student network) and Rec Head for reconstruction (teacher network). ${L}_{hint}$ is calculated from the last two layers of both networks.
• Figure 4: Comparisons with SOTA one-shot MIS methods on BCV.
• Figure 5: Comparison of our method with SOTA one-shot MIS methods on arbitrary cases. Yellow boxes indicate regions where our method is superior.