One-shot segmentation of novel white matter tracts via extensive data augmentation and adaptive knowledge transfer.

Abstract

The use of convolutional neural networks (CNNs) has allowed accurate white matter (WM) tract segmentation on diffusion magnetic resonance imaging (dMRI). To train the CNN-based segmentation models, a large number of scans on which WM tracts are annotated need to be collected, and these annotated scans can be accumulated over a long period of time. However, when novel WM tracts that are different from existing annotated WM tracts are of interest, additional annotations are required for their segmentation. Due to the cost of manual annotations, methods have been developed for few-shot segmentation of novel WM tracts, where the segmentation knowledge is transferred from existing WM tracts to novel WM tracts and the amount of annotated data for novel WM tracts is reduced. Despite these developments, it is desirable to further reduce the amount of annotated data to the one-shot setting with a single annotated image. To address this problem, we develop an approach to one-shot segmentation of novel WM tracts. Our method follows the existing pretraining/fine-tuning framework that transfers segmentation knowledge from existing to novel WM tracts. First, as there is extremely scarce annotated data in the one-shot setting, we design several different data augmentation strategies so that extensive data augmentation can be performed to obtain extra synthetic training data. The data augmentation strategies are based on image masking and thus applicable to the one-shot setting. Second, to address overfitting and knowledge forgetting in the fine-tuning stage that can be more severe given limited training data, we propose an adaptive knowledge transfer strategy that selects the network weights to be updated. The data augmentation and adaptive knowledge transfer strategies are combined to train the segmentation model. Considering that the different data augmentation strategies can generate synthetic data that contain potentially conflicting information, we apply the data augmentation strategies separately, each leading to a different segmentation model. The results predicted by the different models are fused to produce the final segmentation. We validated our method on two brain dMRI datasets, including a public dataset and an in-house dataset. Different settings were considered for the validation, and the results show that the proposed method improves the one-shot segmentation of novel WM tracts.