Combining generative modelling and semi-supervised domain adaptation for whole heart cardiovascular magnetic resonance angiography segmentation

Accurate quantification of whole heart anatomy is required for patient diagnosis and prognosis as well as the evaluation of treatment. Non-invasive medical imaging techniques such as computed tomography angiography (CTA) or cardiovascular magnetic resonance angiography (CMRA) can be used to quantify 3D heart anatomy [16, 26]. Deep Learning techniques can give highly accurate segmentations given sufficient ground truth labels [2, 6, 25] even in the presence of low contrast and high noise images [12, 15]. Despite rapid advancements, a major unsolved problem is the poor adaptability of these methods to different imaging modalities, scanners and acquisition protocols, and the related need for large amounts of labelled data in each domain. Deep learning models trained on one domain do not generalize well to a different target domain [1] with zero or very few labelled cases. The difference in distribution between a large, labelled source domain and an unlabeled target domain is called “domain gap" [10, 22].

Here, we investigate how to bridge the domain gap between CTA images used as source and CMRA images as target. Given sufficient manual domain-specific ground truth labelled data, the UNet architecture had provided state-of-the-art performance in whole heart segmentation applications [25]. In particular, the nn-UNet package provides a self-configuring solution with a range of data augmentation and ensembling tools [11]. However, these models typically do not generalize to other domains. Unsupervised domain adaptation (UDA) methods seek to transfer segmentation ability acquired in one domain to the other without the need for ground truth labels in the target domain. Most UDA methods use generative modelling, in which the domain gap is reduced by generating target domain images from source domain images. Two powerful UDA generative modelling methods are variational autoencoders (VAEs) and generative adversarial networks (GANs). VAEs use probabilistic encoder and decoder networks, optimized using maximum likelihood. Conversely, GANs use generator and discriminator networks which are optimized in an adversarial manner. It is currently unclear which architecture is best for cardiac CMRA applications. Two methods which have provided state-of-the-art performance in whole heart segmentation UDA applications are the synergistic image and feature alignment (SIFA) GAN architecture [7, 8] and the variational approximation for domain adaptation (VARDA) VAE architecture [20] which have both shown greater success in CMRA (source) to CTA (target) domain transfer than vice-versa.

In this paper we improve our preliminary work [14] investigating semi-supervised performance of these two methods, by removing the requirement for pre-processing registration between CTA and CMRA datasets, incorporating a fivefold cross validation for a full statistical analysis, testing a range of supervision levels from 0 to 30%, testing performance on two different CMRA protocols, and comparing with the supervised nn-UNet method. This study is designed to be applied to multi-domain data with a variable number of supervised cases, and to provide a guideline for choosing the best approach for this challenging problem.

All experiments were evaluated through 3D Dice and ASD (mm) metrics. The results were aggregated over all test folds (Additional file 1: Tables S1, S2).

We first provide an overview of the experiments in Figs. 3, 4. These show the effect of adding supervised target cases in the HRMRA and MMWHS datasets respectively. As expected, the performance of GAN, VAE and UNet were higher on the HRMRA dataset than on the MMWHS dataset. This is likely due to the availability of more cases for training, and the higher quality of the HRMRA images. ASD results showed similar patterns to Dice (Additional file 1: Tables S1 and S2). Comparing different levels of supervision per each label, the difference between no supervision and 30% supervision was almost always significant. This was true for all experiments, except for some labels in the UNet method, and for the RA/AO labels in the GAN-based results for the HRMRA dataset. A significant improvement was often found when going from no supervision to 10% supervision, especially in the HRMRA dataset, while it becomes rarer between 10% supervision and 20% supervision, and even more between 20 and 30% experiments.

Figure S1, S2 in Additional file 1 compare performance between methods. Both GAN and VAE significantly outperformed Unet for most labels and supervision levels. In the MMWHS dataset, GAN and VAE methods were not significantly different for LV, LVM and RV labels at 0%, 10%, 20% and 30% supervised target cases. However, GAN outperformed VAE for LA at 0% and 20%, RA at 0% and 10%, AO at 0%, 20% and 30% supervision (\(p<0.05\)). In the HRMRA dataset, GAN outperformed VAE for most labels at every level of supervision.

Figures 5, 6, 7, 8 show examples of the outputs from the GAN-based method and the VAE-based method on each dataset. The first two rows show the images generated by each experiment for CTA or CMRA inputs, the last row shows the ground truth MR segmentation on the left and the predicted ones on the right. The difference in the two generative approaches can be observed. The GAN-based network tries to generate an CMRA from CTA, and vice versa, while the VAE generates less domain-dependent images due to the alignment of the joint latent space. From a comparison between Figs. 5 and 6, and between Figs. 7 and 8, it is also clearly visible that the MMWHS dataset represents a bigger challenge with greater impact of the supervision on the quality of the predicted labels. The difference is more subtle for the HRMRA dataset (Figs. 6 and 8), where less improvement in predicted segmentation quality is seen with increasing supervision, although still present (cf. Fig. 3 and Additional file 1: Table S2).

To validate the consistency across slices, in Fig. 9 we show the 3D label map obtained by each method. Again, this confirms that the HRMRA dataset is easier to segment, and that, even with a 30% supervision, a GAN-based approach is preferable to the fully supervised one. The difference between 1B and 1D for the MMWHS dataset appears to be visually significant, with very scarce labels and poor boundaries drawn by a UNet and a much better result in case D.

Finally, Table 1 summarizes the average signed differences and root-mean-squared error (RMSE) between HRMRA volumes obtained by the ground truth label maps and the ones calculated from outputs of the trained models. This shows that the calculated volumes are substantially closer to the ground truth values in semi-supervised experiments, rather than in unsupervised ones. The corresponding table for MMWHS data is found in Additional file 1: Table S3.

Our results show that adding a small number of supervised cases to a generative modelling domain adaptation method can significantly boost segmentation quality. In our experiments, using two CMRA datasets with different resolution quality, as little as 10% supervision was enough for a significant change from a completely unsupervised approach (Figs. 3, 4). In a situation where very little ground truth is available in one domain, we show that supervised techniques, although optimally trained, are outperformed by generative domain adaptation methods. In the MMWHS dataset, we achieved Dice scores of 0.86 (LV), 0.69 (LVM), 0.74 (RV), 0.73 (LA), 0.79 (RA), 0.74 (AO). This is on average a 12% increase on the original GAN-based approach [8], which reported 0.79 (LV), 0.47 (LVM), 0.62 (LA), 0.65 (AO), and a 10% increase on the results from [20]: 0.74 (LV), 0.47 (LVM), 0.73 (RV), 0.63 (LA), 0.71 (RA) in the original VAE approach. Moreover, in the HRMRA dataset, providing higher resolution in the CMRA domain, we obtained an average Dice of 0.86 across all labels. This compares well with the results of the 2017 Multi-Modality Whole Heart Segmentation (MMWHS) challenge [25], which presented more than 10 algorithms for supervised CTA and CMRA segmentation, with the highest results reporting a maximum 3D Dice of 0.908 \(\pm\) 0.086 for CT and 0.874 \(\pm\) 0.039 for MR. Compared with interobserver errors from a multi-core-lab study [19], the mean and standard deviation of volume differences between core labs was typically 20 ± 10 mL respectively, compared with − 1.8 ± 8.1 mL + for LV and 2.2 ± 10.3 ml for RV volume (Table 1, 30% supervised GAN). Scan-rescan coefficient of variation for LV mass was 5%, or 7.5 g, in [3] compared with 9.5 g for the 30% supervised GAN (Table 1). Therefore, although our results need to be improved to achieve super-human performance, current results are clinically applicable.

In our previous work [14], the training of the generative modelling techniques was limited by a manual registration step, whereas here we focus on boosting the results by relying solely on the original images and the architecture. Important factors that improved these models included: scaling of input images to the GAN architectures in a specific range and training on random slices rather than consecutive ones extracted from a volume, cropping images around the heart for the VAE approach. The GAN approach, once stably trained, gave generally better results, and can generate fake images from both source and target domains (Figs. 5, 6). This can be very useful in applications where a certain anatomy is required in a specific modality, or, in the future, for a fusion between imaging modalities [18]. On the other hand, the VAE approach was easier to train and twice as fast (~ 12 h for single fold training vs ~ 24 h for GAN-based method) and can still output high-quality segmentations. Although we solely optimise and validate our methods based on segmentation quality, we believe the difference in reconstruction quality to be a significant factor. VAE based methods tend to produce smoother images than GANs, since they are maximum likelihood estimations rather than a Nash equilibrium. Moreover, VARDA’s architecture specifically aims at learning features that are domain-invariant during the training stage, hence the reconstructed images are intended to lose most of the features which would make them look realistic (cf. Figure 6 in [20]).

In addition to Fig. 9, in Additional files 2 and 3 we provide videos which show the 3D consistency of the label maps obtained on this dataset by both generative techniques. Indeed, the results achieved and the quality of the HRMRA data suggest that one could move further by performing an accurate segmentation of the smaller structures (arteries, veins) and thus, generating another interesting application directed at patients in need of cardiac catheterisation [26]. Both the GAN- and VAE- based approach could also be modified to directly work on 3D inputs and tackle 4D challenges such as the segmentation of 3D cardiac CINE MRI, investigating motion and geometry of the heart.

Although in this paper we overcome the previous requirement of multi-modality registration [14], the experiments presented here heavily rely on pre-processing steps which include cropping and centering of the input images. These are challenging to identify and widely vary depending on the technique used. They also appear to significantly affect segmentation performance. Indeed, the cropping step seems to be an essential factor for both networks, and particularly in the VAE-based approach, which requires a very close field of view. In our experiments, this was obtained by identifying the centroid of the heart using the ground-truth labels but in clinical practice the performance would highly benefit from a pre-processing bounding box location step. The generalizability of deep learning methods for segmentation would benefit from increased dataset size and variability in terms of MR scanner, sequence and presence of pathologies. Lastly, we limited evaluation to 6 labels, however, extension to pulmonary veins and arteries is possible in future work since the labels are available in the ground truth dataset.

In this study, we compared semi-supervised methods based on generative modelling with a state-of-the-art fully supervised one for the task of CMRA segmentation across two datasets with different original resolution. We demonstrated that, in absence of many ground truth cases, a domain adaptation approach is beneficial, and this can be used to accurately segment bigger structures as well as minor ones, and to generate synthetic images of specific imaging modalities.