Adaptive wavelet-VNet for single-sample test time adaptation in medical image segmentation.

Abstract

In medical image segmentation, a domain gap often exists between training and testing datasets due to different scanners or imaging protocols, which leads to performance degradation in deep learning-based segmentation models. Given the high cost of manual labeling and the need for privacy protection, it is often challenging to annotate the testing (target) domain data for model fine-tuning or to collect data from different domains to train domain generalization models. Therefore, using only unlabeled target domain data for test-time adaptation (TTA) presents a more practical but challenging solution. To improve the segmentation accuracy of deep learning-based models on unseen datasets, and especially to enhance the efficiency and stability of TTA for individual samples from heterogeneous domains. In this study, we proposed to dynamically adapt a wavelet-VNet (WaVNet) to unseen target domains with a hybrid objective function, based on each unlabeled test sample during the test time. We embedded multiscale wavelet coefficients into a V-Net encoder and adaptively adjusted the spatial and spectral features according to the input, and the model parameters were optimized by three loss functions. We integrated a shape-aware loss to focus on the foreground segmentations, a Refine loss to correct the incomplete and noisy segmentations caused by domain shifts, and an entropy loss to promote the global consistency of the segmentations. We evaluated the proposed method on multidomain liver and prostate segmentation datasets to assess its advantages over other TTA methods. For the source domain model training of the liver dataset, we used 15 3D MR image samples for training and 5 for validation. Correspondingly, for the prostate dataset, we used 22 3D MR image samples for training and 7 for validation. In the target domain, we used a single 3D MR image sample for adaptation and testing. The total number of testing samples is 60 in the liver dataset (for 3 different domains) and 116 in the prostate dataset (for 6 different domains). The proposed method showed the highest segmentation accuracy among all methods, achieving a mean(±SD) Dice coefficient (DSC) of 78.10±5.23% and a mean 95th Hausdorff distance (HD95) of 15.52±5.84mm on the liver dataset; and a mean DSC of 80.02±3.89% and a mean HD95 of 9.18±3.47mm on the prostate dataset. The DSC is 11.67% (in absolute terms) and 15.27% higher than that of the baseline (no adaptation) method, for the liver and the prostate datasets, respectively. The proposed adaptive WaVNet enhanced the image segmentation accuracy from unseen domains during the test time via unsupervised learning and multi-objective optimization. It can benefit clinical applications where data are scarce or with changing data distributions, including online adaptive radiotherapy. The code will be released at: https://github.com/sanny1226/WaVNet.