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

2024-08-07

arXiv.org

AI Summary

In single-label medical image segmentation, performance degradation arises from label propagation errors and poor synthetic image quality. To address this, we propose a reconstruction-guided knowledge distillation framework. Our method employs registration-based augmentation to generate high-fidelity synthetic labeled samples and establishes a teacher–student collaborative learning paradigm, wherein reconstructed images serve as soft supervision signals for feature distillation. This marks the first integration of image reconstruction as a guidance source in knowledge distillation—effectively circumventing registration inaccuracies and interference from low-quality synthesis—without requiring explicit image registration or aggressive data augmentation. Evaluated on three cross-modality benchmarks—OASIS (T1-MRI), BCV (abdominal CT), and VerSe (vertebral CT)—our approach significantly outperforms state-of-the-art one-shot segmentation methods, achieving superior segmentation accuracy and generalization capability.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Multi-class/Multi-label Learning & Extreme Classification
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs

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.

Problem

• Medical Image Segmentation
• Atlas Label Propagation Error
• Synthetic Image Quality

Innovation

• Medical Image Segmentation
• Registration Technique
• Student-Teacher Model