Gaussian Process Emulators for Few-Shot Segmentation in Cardiac MRI

Bruno Viti,Franz Thaler,Kathrin Lisa Kapper,Martin Urschler,Martin Holler,Elias Karabelas
2024-11-11
Abstract:Segmentation of cardiac magnetic resonance images (MRI) is crucial for the analysis and assessment of cardiac function, helping to diagnose and treat various cardiovascular diseases. Most recent techniques rely on deep learning and usually require an extensive amount of labeled data. To overcome this problem, few-shot learning has the capability of reducing data dependency on labeled data. In this work, we introduce a new method that merges few-shot learning with a U-Net architecture and Gaussian Process Emulators (GPEs), enhancing data integration from a support set for improved performance. GPEs are trained to learn the relation between the support images and the corresponding masks in latent space, facilitating the segmentation of unseen query images given only a small labeled support set at inference. We test our model with the M&Ms-2 public dataset to assess its ability to segment the heart in cardiac magnetic resonance imaging from different orientations, and compare it with state-of-the-art unsupervised and few-shot methods. Our architecture shows higher DICE coefficients compared to these methods, especially in the more challenging setups where the size of the support set is considerably small.
Computer Vision and Pattern Recognition,Artificial Intelligence,Machine Learning
What problem does this paper attempt to address?
This paper attempts to address the problem of insufficient data annotation in cardiac magnetic resonance imaging (MRI) segmentation, especially when training and test images come from different distributions (e.g., different cardiac slice orientations). Specifically, the goals of the paper are: 1. **Reduce the dependence on a large amount of annotated data**: Traditional deep - learning methods usually require a large amount of annotated data for training, which is a challenge in the medical field, especially in cardiac MRI segmentation, because high - quality annotated data is very scarce. 2. **Improve the adaptability of the model to different cardiac image orientations**: Many existing methods perform well when dealing with images of different modalities or sources, but have limitations in adapting to new image orientations (such as from the short - axis to the long - axis). This research aims to overcome this problem by introducing few - shot learning (FSS). 3. **Propose a new method combining U - Net and Gaussian Process Emulators (GPEs)**: By combining FSS with the U - Net architecture and using GPEs to extract information from a small support set, efficient segmentation of unseen query images can be achieved. ### Specific Problem Description The paper explores the following key issues: - **How to achieve efficient cardiac MRI segmentation with limited annotated data**: By introducing a few - shot learning framework, the segmentation task of new images can be completed with only a small annotated support set. - **How to handle image segmentation of different cardiac slice orientations**: Especially during the transition from the short - axis (SA) to the long - axis (LA), ensure that the model can effectively generalize and maintain high segmentation accuracy. - **How to integrate the information in the support set to improve the segmentation performance of query images**: By learning the relationship between support images and their corresponding masks through GPEs, the segmentation of query images can be better represented and supported in the latent space. ### Method Overview To achieve the above goals, the paper proposes a few - shot segmentation model based on U - Net and GPEs. The main steps include: 1. **Encoding stage**: Use the contraction path of U - Net to perform feature extraction on query images and support set images, and map them to the latent space. 2. **GPEs training**: Train GPEs in the latent space to learn the relationship between support image features and their corresponding masks. 3. **Prediction stage**: Given the features of new query images, use GPEs to infer their corresponding masks and pass this information to the expansion path of U - Net to generate the final segmentation result. In this way, this method can effectively segment cardiac MRI images in different orientations with only a small amount of annotated data, and shows better performance than existing methods in experiments. ### Experimental Verification The paper uses the M&Ms - 2 public dataset for evaluation to verify the effectiveness of this method in the short - axis - to - long - axis segmentation task. The experimental results show that this method achieves significantly better performance than other methods in the 1 - shot setting, and the segmentation accuracy further improves as the size of the support set increases. ### Summary This research successfully solves the problems of insufficient data annotation and cross - orientation generalization in cardiac MRI segmentation by combining U - Net and GPEs, providing an effective solution for few - shot medical image segmentation.