A transfer learning approach to few-shot segmentation of novel white matter tracts.

Abstract

Convolutional neural networks(CNNs) have achieved state-of-the-art performance for white matter(WM) tract segmentation based on diffusion magnetic resonance imaging(dMRI). The training of the CNN-based segmentation model generally requires a large number of manual delineations of WM tracts, which can be expensive and time-consuming. Although it is possible to carefully curate abundant training data for a set of WM tracts of interest, there can also be novel WM tracts-i.e., WM tracts that are not included in the existing annotated WM tracts-that are specific to a new scientific problem, and it is desired that the novel WM tracts can be segmented without repeating the laborious collection of a large number of manual delineations for these tracts. One possible solution to the problem is to transfer the knowledge learned for segmenting existing WM tracts to the segmentation of novel WM tracts with a fine-tuning strategy, where a CNN pretrained for segmenting existing WM tracts is fine-tuned with only a few annotated scans to segment the novel WM tracts. However, in classic fine-tuning, the information in the last task-specific layer for segmenting existing WM tracts is completely discarded. In this work, based on the pretraining and fine-tuning framework, we propose an improved transfer learning approach to the segmentation of novel WM tracts in the few-shot setting, where all knowledge in the pretrained model is incorporated into the fine-tuning procedure. Specifically, from the weights of the pretrained task-specific layer for segmenting existing WM tracts, we derive a better initialization of the last task-specific layer for the target model that segments novel WM tracts. In addition, to allow further improvement of the initialization of the last layer and thus the segmentation performance in the few-shot setting, we develop a simple yet effective data augmentation strategy that generates synthetic annotated images with tract-aware image mixing. To validate the proposed method, we performed experiments on brain dMRI scans from public and private datasets under various experimental settings, and the results indicate that our method improves the performance of few-shot segmentation of novel WM tracts.