Generative AI Enables Medical Image Segmentation in Ultra Low-Data Regimes

Li Zhang,Basu Jindal,Ahmed Alaa,Robert Weinreb,David Wilson,Eran Segal,James Zou,Pengtao Xie

2024-08-31

Abstract:Semantic segmentation of medical images is pivotal in applications like disease diagnosis and treatment planning. While deep learning has excelled in automating this task, a major hurdle is the need for numerous annotated segmentation masks, which are resource-intensive to produce due to the required expertise and time. This scenario often leads to ultra low-data regimes, where annotated images are extremely limited, posing significant challenges for the generalization of conventional deep learning methods on test images. To address this, we introduce a generative deep learning framework, which uniquely generates high-quality paired segmentation masks and medical images, serving as auxiliary data for training robust models in data-scarce environments. Unlike traditional generative models that treat data generation and segmentation model training as separate processes, our method employs multi-level optimization for end-to-end data generation. This approach allows segmentation performance to directly influence the data generation process, ensuring that the generated data is specifically tailored to enhance the performance of the segmentation model. Our method demonstrated strong generalization performance across 9 diverse medical image segmentation tasks and on 16 datasets, in ultra-low data regimes, spanning various diseases, organs, and imaging modalities. When applied to various segmentation models, it achieved performance improvements of 10-20\% (absolute), in both same-domain and out-of-domain scenarios. Notably, it requires 8 to 20 times less training data than existing methods to achieve comparable results. This advancement significantly improves the feasibility and cost-effectiveness of applying deep learning in medical imaging, particularly in scenarios with limited data availability.

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper aims to address the challenges of medical image segmentation in ultra-low data environments. Specifically, although deep learning performs excellently in automatically executing medical image segmentation tasks, it requires a large number of annotated segmentation masks, which is resource-intensive in practice. This resource-intensive demand leads to "ultra-low data" scenarios, where annotated images are extremely limited, posing significant challenges to the generalization performance of traditional deep learning methods on test images.

To address this issue, the authors propose a generative deep learning framework—GenSeg. The uniqueness of this framework lies in its ability to achieve end-to-end data generation through multi-level optimization, thereby generating high-quality paired segmentation masks and medical images as auxiliary data to train robust models. This approach allows segmentation performance to directly influence the data generation process, ensuring that the generated data is specifically aimed at enhancing the performance of the segmentation model. Experimental results show that GenSeg has significant advantages over existing methods in multiple medical image segmentation tasks and can achieve comparable performance with only a small number of training samples. Additionally, GenSeg demonstrates excellent generalization capabilities in segmentation tasks across different domains (both same-domain and cross-domain).