One-Shot Segmentation of Novel White Matter Tracts via Extensive Data Augmentation

Suppose we are given a CNN-based model $\mathcal{M}_{\rm{e}}$ that segments $M$ existing WM tracts, for which a large number of annotations have been accumulated for training. We are interested in training a CNN-based model $\mathcal{M}_{\rm{n}}$ for segmenting $N$ novel WM tracts, for which only one scan is annotated due to annotation cost.

Existing work [10] has attempted to address this problem with a transfer learning strategy based on the pretraining and fine-tuning framework, where $\mathcal{M}_{\rm{e}}$ and $\mathcal{M}_{\rm{n}}$ share the same network structure for feature extraction, and their last task-specific layers are different. In classic fine-tuning, the network weights of the learned feature extraction layers of $\mathcal{M}_{\rm{e}}$ are used to initialize the feature extraction layers of $\mathcal{M}_{\rm{n}}$, and the task-specific layer of $\mathcal{M}_{\rm{n}}$ is randomly initialized. Then, all weights of $\mathcal{M}_{\rm{n}}$ are fine-tuned with the single scan annotated for novel WM tracts. However, the classic fine-tuning strategy discards the information in the task-specific layer of $\mathcal{M}_{\rm{e}}$. As different WM tracts cross or overlap, existing and novel WM tracts can be correlated, and the task-specific layer of $\mathcal{M}_{\rm{e}}$ for segmenting existing WM tracts may also bear relevant information for segmenting novel WM tracts. Therefore, to exploit all the knowledge learned in $\mathcal{M}_{\rm{e}}$, in [10] an improved fine-tuning strategy is developed, which, after derivation, can be conveniently achieved with a warmup stage. Specifically, the feature extraction layers of $\mathcal{M}_{\rm{n}}$ are first initialized with those of $\mathcal{M}_{\rm{e}}$. Then, in the warmup stage, the feature extraction layers of $\mathcal{M}_{\rm{n}}$ are fixed and only the last task-specific layer of $\mathcal{M}_{\rm{n}}$ (randomly initialized) is learned with the single annotated image. Finally, all weights of $\mathcal{M}_{\rm{n}}$ are jointly fine-tuned with the single annotated image.

Although the transfer learning approach in [10] has improved the few-shot segmentation of novel WM tracts, when the training data for novel WM tracts is extremely scarce with only one annotated image, the segmentation is still challenging. Therefore, we continue to explore the problem of one-shot segmentation of novel WM tracts. Based on the pretraining and fine-tuning framework developed in [10], we propose to more effectively exploit the information in the single annotated image with extensive data augmentation for network training. Suppose the annotated image is $\mathbf{X}$ and its annotation is $\mathbf{Y}$ (0 for background and 1 for foreground); then we obtain a set of synthetic annotated training images $\widetilde{\mathbf{X}}$ and the synthetic annotations $\widetilde{\mathbf{Y}}$ by transforming $\mathbf{X}$ and $\mathbf{Y}$. We develop several data augmentation strategies for the purpose, which are described below.

We use the state-of-the-art TractSeg architecture [17] for volumetric WM tract segmentation as our backbone segmentation network, which takes fiber orientation maps as input. Like [17], the fiber orientation maps are computed with constrained spherical deconvolution (CSD) [15] or multi-shell multi-tissue CSD (MSMT-CSD) [7] for single-shell or multi-shell dMRI data, respectively, and three fiber orientations are allowed in the network input [17]. We follow [17] and perform 2D WM tract segmentation for each image view separately, and the results are fused to obtain the final 3D WM tract segmentation.

The proposed data augmentation is performed offline. Since given $N$ novel WM tracts TC1 or TC2 can produce at most $2^{N}-1$ different images, we set the number of synthetic scans produced by each data augmentation strategy to $\min(2^{N}-1,100)$. Note that traditional data augmentation, such as elastic deformation, scaling, intensity perturbation, etc., is applied online in TractSeg to training images. Thus, these operations are also applied to the synthetic training data online. The training configurations are set according to TractSeg, where Adamax [8] is used to minimize the binary cross entropy loss with a batch size of 56, an initial learning rate of 0.001, and 300 epochs [17]. The model corresponding to the epoch with the best segmentation accuracy on a validation set is selected.

For evaluation, experiments were performed on the publicly available Human Connectome Project (HCP) dataset [16] and an in-house dataset. The dMRI scans in the HCP dataset were acquired with 270 diffusion gradients (three $b$-values) and a voxel size of 1.25 mm isotropic. In [17] 72 major WM tracts were annotated for the HCP dataset, and the annotations are also publicly available. For the list of the 72 WM tracts, we refer readers to [17]. The dMRI scans in the in-house dataset were acquired with 270 diffusion gradients (three $b$-values) and a voxel size of 1.7 mm isotropic. In this dataset, only ten of the 72 major WM tracts were annotated due to the annotation cost.

Following [10], we selected the same 60 WM tracts as existing WM tracts, and a segmentation model was pretrained for these tracts with the HCP dMRI scans, where 48 and 15 annotated scans were used as the training set and validation set, respectively. To evaluate the performance of one-shot segmentation of novel WM tracts, we considered a more realistic and challenging scenario where novel WM tracts are to be segmented on dMRI scans that are acquired differently from the dMRI scans annotated for existing WM tracts. Specifically, instead of using the original HCP dMRI scans for segmenting novel WM tracts, we generated clinical quality scans from them. The clinical quality scans were generated by selecting only 34 diffusion gradients at $b$ = 1000 s/mm² and downsampling the selected diffusion weighted images in the spatial domain by a factor of two to the voxel size of 2.5 mm isotropic. The tract annotations were also downsampled accordingly. Since the dMRI scans in the in-house dataset were acquired differently from the original HCP dMRI scans, they were directly used for evaluation together with their original annotations.

The two datasets were used to evaluate the accuracy of one-shot segmentation of novel WM tracts separately based on the model pretrained on the original HCP dataset for segmenting existing WM tracts. We considered three cases of novel WM tracts for each dataset. For the clinical quality scans, the three cases are referred to as CQ1, CQ2, and CQ3, and for the in-house dataset, the three cases are referred to as IH1, IH2, and IH3. The details about these cases are summarized in Table 1. For each dataset, only one scan was selected from it for network training, together with the corresponding annotation of novel WM tracts.²²2The single annotated scan was also used as the validation set for model selection. For the clinical quality dataset and the in-house dataset, 30 and 16 scans were used for testing, respectively, where the annotations of novel WM tracts were available and only used to measure the segmentation accuracy.

We compared the proposed method with two competing methods, which are the classic fine-tuning strategy and the improved fine-tuning strategy [10] described in Section 2.1 with the same pretrained model and single annotated training scan used by the proposed method. Both competing methods were integrated with TractSeg [17] like the proposed method. For convenience, the classic fine-tuning strategy and the improved fine-tuning strategy are referred to as CFT and IFT, respectively. Note that as shown in [10], in the one-shot setting directly training a model that segments novel WM tracts from scratch without the pretrained model would lead to segmentation failure. Thus, this strategy was not considered.

We first qualitatively evaluated the proposed method. Examples of the segmentation results for novel WM tracts are shown in Fig. 1 for the proposed and competing methods, where the annotations are also displayed for reference. For demonstration, here we show the results obtained in the case of CQ1 for CST_left and OR_right. We can see that the segmentation results of our method better resemble the annotations than those of the competing methods.

Next, we quantitatively evaluated our method. The Dice coefficient was computed for the segmentation result of each novel WM tract on each test scan for each case. For demonstration, we have listed the average Dice coefficient of each novel WM tract in Table 2 for the cases of CQ1 and IH1. For each tract and each case, our method has a higher average Dice coefficient than the competing methods, and the improvement is statistically significant. We have also summarized the mean of the average Dice coefficients of the novel WM tracts for all the cases in Table 3 (upper half table). In all cases our method outperforms the competing methods with higher mean Dice coefficients.

Finally, we confirmed the benefit of each proposed data augmentation strategy, as well as the benefit of ensembling their results. For each case, the mean value of the average Dice coefficients of all novel WM tracts was computed for the segmentation results achieved with RC1, RC2, TC1, or TC2 individually, and the results are also given in Table 3 (lower half table). Compared with the results of IFT that did not use the proposed data augmentation, the integration of IFT with RC1, RC2, TC1, or TC2 led to improved segmentation accuracy, which indicates the individual benefit of each proposed data augmentation strategy. In addition, the Dice coefficients of the proposed method achieved with ensembling are higher than those achieved with a single data augmentation strategy, which confirms the benefit of ensembling. Note that there is not a data augmentation strategy that is better or worse than the others in all cases, which is possibly because of the randomness in RC1 and RC2 and the dependence of TC1 and TC2 on the spatial coverages of the novel WM tracts that vary in different cases.