Generative augmentations for improved cardiac ultrasound segmentation using diffusion models

TL;DR

This study tackles the generalization gap in cardiac ultrasound segmentation caused by scarce labeled data and varying annotation conventions. It introduces unconditional diffusion-model generative augmentations, implemented via RePaint, to diversify training data without altering model architectures. Across HUNT4 and CAMUS datasets, the approach yields realistic images and improved segmentation robustness, including better cross-domain performance and tighter agreement in automatic EF estimation. The method demonstrates practical impact by enhancing generalizability in clinical settings and is released as an open-source library for broad adoption. Overall, diffusion-based data augmentation emerges as a powerful avenue to bolster segmentation reliability in echocardiography under limited labeling resources.

Abstract

One of the main challenges in current research on segmentation in cardiac ultrasound is the lack of large and varied labeled datasets and the differences in annotation conventions between datasets. This makes it difficult to design robust segmentation models that generalize well to external datasets. This work utilizes diffusion models to create generative augmentations that can significantly improve diversity of the dataset and thus the generalisability of segmentation models without the need for more annotated data. The augmentations are applied in addition to regular augmentations. A visual test survey showed that experts cannot clearly distinguish between real and fully generated images. Using the proposed generative augmentations, segmentation robustness was increased when training on an internal dataset and testing on an external dataset with an improvement of over 20 millimeters in Hausdorff distance. Additionally, the limits of agreement for automatic ejection fraction estimation improved by up to 20% of absolute ejection fraction value on out of distribution cases. These improvements come exclusively from the increased variation of the training data using the generative augmentations, without modifying the underlying machine learning model. The augmentation tool is available as an open source Python library at https://github.com/GillesVanDeVyver/EchoGAINS.

Figures (16)

• Figure 1: Examples of the generative augmentations types used in this work. All the examples are generated from the same original image shown in the top-left corner.
• Figure 2: Heatmaps of pixels belonging to the LV after resizing to $256\times256$. This illustrates the difference in scan depth variation and LV positioning in the two datasets, before and after generative augmentations.
• Figure 3: Example frames from the CAMUS and HUNT4 datasets. The CAMUS and HUNT4 dataset contain the same cardiac views, but the frames in HUNT4 are consistently LV-focused, while those in CAMUS are not. The annotations conventions are also different in both datasets, which can be seen clearly in the thickness of the annotated myocardium (blue).
• Figure 4: Distribution of imaging depths in CAMUS and HUNT4.
• Figure 5: Distribution of sector angles in CAMUS and HUNT4.