Cross-Domain Distribution Alignment for Segmentation of Private Unannotated 3D Medical Images

TL;DR

A new source-free Unsupervised Domain Adaptation (UDA) method based on estimating the internally learned distribution of a relevant source domain by a base model and then generating pseudo-labels that are used for enhancing the model refinement through self-training is introduced.

Abstract

Manual annotation of 3D medical images for segmentation tasks is tedious and time-consuming. Moreover, data privacy limits the applicability of crowd sourcing to perform data annotation in medical domains. As a result, training deep neural networks for medical image segmentation can be challenging. We introduce a new source-free Unsupervised Domain Adaptation (UDA) method to address this problem. Our idea is based on estimating the internally learned distribution of a relevant source domain by a base model and then generating pseudo-labels that are used for enhancing the model refinement through self-training. We demonstrate that our approach leads to SOTA performance on a real-world 3D medical dataset.

Figures (6)

• Figure 1: The proposed Model Adaptation algorithm for 3D medical image segmentation. (top) the workflow begins with pre-training the segmentation model in a supervised setting using the source domain annotated samples. (bottom, left) to preserve privacy, the learned source distribution at the output-space of the encoder is estimated using a GMM. The GMM represents the internal feature distributions and is shared for model adaptation. (bottom, right) the target domain images undergo patch extraction and feature encoding. The purple encoder-decoder pair indicates a model initialized with parameters from the pre-trained source model. This model is adapted in two steps: (i) log-likelihood calculation compares the target domain GMM against the source domain GMM to find the most similar feature distributions; (ii) shape estimation involves feeding the target domain data into the frozen source-trained model to determine pixel-to-label assignments based on the source-trained parameters. These assignments provide an estimated shape for the target domain data, allowing for estimating the GMM latent features to be positioned close to these estimated shapes. The Sliced Wasserstein Distance is then minimized as the loss function to align the target domain feature distribution with the the source domain distribution through the estimated GMM, enhancing segmentation quality.
• Figure 2: Qualitative performance: examples of the segmented frames for MMWHS dataset. From top to bottom in each case: input images, supervised learning predictions, source-trained model predictions, predictions based on our method, and ground truth provided by radiologists.
• Figure 3: UMAP visualizations for GMM samples: (a) the average for all patients' GMMs; (b) a GMM from one random crop of patient 1; (c) a GMMs from one random crop of patient 2; and (d) a GMMs from one random crop of patient 3. The colors represent different labels: red for MYO, yellow for LAC, green for LVC, and blue for AA. Cropping during adaptation is a benefit for adapted training because the target data crop sometimes contains only part of the labels.
• Figure 4: Comparison of internal latent distributions during post-adaptation: (a) the internal latent distribution after post-adaptation using both shape and count estimation; (b) the internal latent distribution after post-adaptation without shape estimation but using the estimated count of each label; (c) the internal latent distribution after post-adaptation without both shape estimation and estimated count of each label. The colors represent different labels: red for MYO, yellow for LAC, green for LVC, and blue for AA. It can be observed that cropping during adaptation enhances training by focusing on label subsets present in the target data crop.
• Figure 5: Comparison of internal latent distributions during different adaptation phases: (a) the internal latent distribution with input source data before adaptation; (b) the internal latent distribution with input target data before adaptation; (c) the internal latent distribution with input target data after adaptation. The colors represent different labels: red for MYO, yellow for LAC, green for LVC, and blue for AA.