Uncertainty-Guided Efficient Interactive Refinement of Fetal Brain Segmentation from Stacks of MRI Slices

We propose MG-Net to obtain simultaneous initial segmentation and uncertainty estimation in real time with a single network and a single forward pass for inference, which is more efficient for interactive segmentation than typical uncertainty estimation methods including MC dropout [3] and model ensemble [5, 8]. As shown in Fig. 2, we modify U-Net [11] by using Grouped Convolution [7].

An $N$-grouped convolution layer splits the input feature map along the channel dimension into $N$ groups each with $C_{i}$ channels. For each group, it uses a convolution kernel of shape $C_{o}\times C_{i}\times h\times w$ respectively, where $h\times w$ is the spatial size of the kernel. Therefore, $N$ independent feature maps are obtained each with $C_{o}$ channels. They are concatenated into a single one with $N\times~{}C_{o}$ channels. Note that there is no correlation between different groups in the output feature map of an $N$-grouped convolution layer. Similarly, we implement up-sampling layers by transposed $N$-grouped convolutions, and extend standard channel concatenation to group-wise concatenation to keep the $N$ groups independent of each other. Let $F_{1}$ and $F_{2}$ represent two feature maps each with $N$ groups. We first concatenate the $n$-th group of $F_{1}$ with the $n$-th group of $F_{2}$, and denote the result as $\hat{F}_{n}$. Then $\hat{F}_{1}$, $\hat{F}_{2}$, …, and $\hat{F}_{N}$ are concatenated as the group-wise concatenation of $F_{1}$ and $F_{2}$. The group-wise concatenation is used for skip connection between feature maps in the encoder and their counterparts in the decoder, as shown in Fig. 2. At the last layer of the decoder, we apply softmax to each feature group respectively (i.e., group-wise softmax) to obtain $N$ probability predictions. Therefore, MG-Net can be seen as an ensemble of $N$ parallel sub-networks, and they are randomly initialized and trained with dropout to obtain diversity. At the lowest resolution level of MG-Net, we set group number to one to allow communication of these $N$ sub-networks for better performance.

For training, we apply the same segmentation loss function (i.e., Dice loss [10]) to the $N$ predictions respectively and average their loss values for back-propagation. At inference time, we take the pixel-wise mean value and variance of the $N$ probability predictions as the final segmentation probability map and pixel-level uncertainty estimation, respectively. The mean foreground probability map is thresholded by 0.5 to obtain a binary segmentation mask.

For a stack of $M$ slices, the CNN is able to obtain accurate segmentation for most slices and only a few slices may require manual refinement [12]. To avoid unnecessary and time-consuming manual search for mis-segmentations of each slice, we ask the user to check and refine only $M^{\prime}$ ($M^{\prime}<M$) slices with the highest slice-level uncertainty. For a slice $S$, its binary segmentation result $Y$ and pixel-level uncertainty map $U$, a naive slice-level uncertainty can be defined as $\nu^{*}=\sum_{\mathbf{x}}U_{\mathbf{x}}$ where $U_{\mathbf{x}}$ is the uncertainty of pixel $\mathbf{x}$. However, this may lead a slice with a small target to be neglected as it often has a low value of $\nu^{*}$ due to a small uncertain region. To address this problem, we alternatively define the slice-level uncertainty as $\nu=(\sum_{\mathbf{x}}U_{\mathbf{x}})/(\sum_{\mathbf{x}}Y_{\mathbf{x}}+\zeta)$, which is normalized by the segmented region size. $\zeta$ is a small number for numerical stability. We iteratively fetch the slice $S_{m}$ with the $m$-th highest $\nu$ to ask for interactive refinement, where $m$ = 1, 2, …, and the iteration terminates when $m>M^{\prime}$, as shown in Fig. 1. $M^{\prime}$ is a predefined number, such as 60% of $M$ according to the performance of the CNN on the validation set. We also use an early termination strategy when refinement is not needed for three consecutive fetched slices.

For interactive refinement, we use a Distance Regularized Level Set Evolution (DRLSE) [9] due to its efficiency and extend it with an interaction-constraint term, which is named as I-DRLSE. Let $\phi$ denote the level set function that is initialized as the signed distance transform of the initial segmentation result. We define an energy function as $E(\phi)=\alpha E_{r}+\beta E_{u}+\lambda E_{l}+\mu E_{d}$, where $\alpha$, $\beta$, $\lambda$, $\mu$ are weighting parameters. $E_{r}$, $E_{u}$, $E_{l}$ and $E_{d}$ are the region, user-interaction, length and distance regularization terms, respectively. $E_{l}=\int_{\Omega}\delta_{\epsilon}(\phi(\mathbf{x}))|\nabla\phi(\mathbf{x})|d\mathbf{x}$ and $E_{d}=\int_{\Omega}p(|\nabla\phi(\mathbf{x})|)d\mathbf{x}$, where $\delta_{\epsilon}$ and $p()$ are the smoothed Dirac delta function and double-well potential function as in [9] respectively. As the target region in the fetal MRI image has an inhomogeneous appearance that brings challenges to standard intensity-based level set methods, we define the region term based on the foreground probability map $P$ obtained by MG-Net instead of the original image:

where $H_{\epsilon}$ is the smoothed Heaviside function as in [9]. $c_{1}$ and $c_{2}$ are the average foreground probability inside and outside the current level set contour, respectively. Our proposed user-interaction term is:

where $\eta(\mathbf{x})$ is user-interaction-derived likelihood of pixel $\mathbf{x}$ being the foreground. Let $\mathcal{F}$ and $\mathcal{B}$ represent the set of pixels specified as the foreground and background by the user interactions, respectively. Inspired by [16], we use $g_{\mathbf{x}}^{\mathcal{F}}$ to represent the geodesic distance between $\mathbf{x}$ and $\mathcal{F}$, and define $\eta(\mathbf{x})=e^{-G_{\mathbf{x}}^{\mathcal{F}}}/(e^{-G_{\mathbf{x}}^{\mathcal{F}}}+e^{-G_{\mathbf{x}}^{\mathcal{B}}})$, where $G_{\mathbf{x}}^{\mathcal{F}}=\min(g_{\mathbf{x}}^{\mathcal{F}},D)$ and $G_{\mathbf{x}}^{\mathcal{B}}=\min(g_{\mathbf{x}}^{\mathcal{B}},D)$ and $D$ is a threshold value to ensure that only a local region is affected by the interactions. We set $g_{\mathbf{x}}^{\mathcal{F}}$ or $g_{\mathbf{x}}^{\mathcal{B}}$ as $D$ when $\mathcal{F}$ or $\mathcal{B}$ is empty. $E_{u}$ is infinite if the segmentation results conflict with the user interactions.

MRI scans of 35 fetuses in the second trimester were collected by Single Shot Fast Spin Echo (SSFSE) with pixel size 0.74 mm-1.58 mm and inter-slice spacing 3 mm-4 mm. Each fetus had three scans in axial, sagittal and coronal views respectively, leading to 105 stacks in total. We randomly selected 72 stacks from 24 patients, 9 stacks from 3 patients and 24 stacks from 8 patients for training, validation and testing, respectively. Manual segmentations of these images were used as the ground truth.

We implemented the CNNs in Pytorch on a Ubuntu desktop with an NVIDIA GTX 1080 Ti GPU and developed a PyQt GUI for user interactions. I-DRLSE was implemented in Python and it ran on the CPU. To train MG-Net, we used Dice loss [10] and Adam optimizer with weight decay $10^{-5}$, mini-batch of 24 slices and learning rate $10^{-4}$. The training was ended when performance on the validation set stopped to increase for 5k iterations. The group number $N$ in MG-Net was 4 and the channel number parameter $C$ in Fig. 2 was 16. For interactive refinement, $M^{\prime}$ was set to $0.6M$ as refinement was not needed for more than half of the slices in the validation set. For I-DRLSE, the maximal evolution step number was 200 and $\alpha$=0.1, $\beta$=0.5, $\lambda$=0.3, $\mu$=0.005, $D$=4.0 based on grid search on the validation set. The segmentation accuracy was measured by Dice similarity and Average Symmetric Surface Distance (ASSD) between segmentation results and the ground truth.

Our MG-Net was compared with MC dropout [3] with 10 folds, ensemble [8] of 5 models, and a Bayesian network [4] for uncertainty estimation. We implemented all these methods using U-Net [11] as the backbone. To measure ability of the uncertainty to indicate mis-segmentation, we used Uncertainty-Error Overlap (UEO, i.e., Dice as in [5]) and Relative Volume Error (RVE) between thresholded uncertain region and mis-segmented region. The optimal threshold value for each method was determined based on the validation set. Quantitative evaluation results are shown in Table 1. Compared with ensemble of U-Net, our MG-Net obtains comparable segmentation accuracy and higher uncertainty estimation quality. MG-Net takes 0.29s in average for simultaneous automatic segmentation and uncertainty estimation of a stack, which is far more efficient than MC dropout and model ensemble and more suitable for interactive segmentation. The Bayesian network method also obtains uncertainty estimation in a fast speed, but with a reduced segmentation accuracy. Visual comparison in Fig. 3 shows that MG-Net obtains better consistency between mis-segmentation and uncertain regions than the other uncertainty estimation methods.

Our I-DRLSE was firstly validated with slice-level refinement, and compared with: 1) CPU-based Graph Cut²²2Code from the maxflow-v3.01 library: https://vision.cs.uwaterloo.ca/code/ as implemented in [15], and 2) training a refinement CNN (a.k.a., R-Net) [16] that takes the initial segmentation and user interactions as input. We implemented the R-Net using U-Net [11] as the backbone, and trained it with simulated interactions on initial segmentation obtained by MG-Net following [16]. As not all the slices in a stack require refinement, we randomly selected 100 obviously mis-segmented slices from the test set, and used the same set of user interactions on the initial segmentation for comparison. Fig. 4(a) shows two cases where over- and under-segmentation exist in the initial segmentation respectively, which demonstrates that I-DRLSE obtains higher refinement accuracy than Graph Cut and R-Net with the same initial segmentation and interactions. Fig. 4(b) shows a quantitative comparison of these refinement methods. It can be observed that I-DRLSE leads to higher Dice scores and lower ASSD values than Graph Cut and R-Net for refinement. The average slice-level machine time for CPU-based I-DRLSE was 0.45s, which is slower than that of Graph Cut (0.07s) and R-Net (0.12s), but still acceptable for fast response of user interactions. The efficiency of I-DRLSE could be further improved by a GPU-optimized implementation in the future.

For stack-level segmentation, our UGIR was compared with two variants: 1) UGIR(-U) denoting that the user manually searches mis-segmentation slice-by-slice to provide interactions (i.e., not guided by uncertainty) for refinement using I-DRLSE; 2) UGIR(*) that denotes using the naive slice-level uncertainty $\nu^{*}$ to guide user interactions. They were also compared with two existing interactive segmentation methods: using Graph Cut [1] for interactive segmentation from scratch, and DeepIGeoS [16] that uses two CNNs for initial segmentation and interactive refinement, respectively. We re-implemented DeepIGeoS following the training method in [16]. One user employed these methods to segment the fetal brain from the testing fetal MRI stacks respectively, where the interactions in a slice could be given multiple times until the result was accepted. Quantitative evaluation results are shown in Table 2. Compared with DeepIGeoS [16], our UGIR obtained similar final accuracy ($p$-value $>$ 0.05 based on a paired t-test), but reduced the runtime from 75.38s to 48.46s. By using uncertainty-guided user interactions, UGIR improved the efficiency by near 30% from UGIR(-U). UGIR(*) and UGIR took almost the same runtime, but UGIR achieved higher accuracy.

This work was supported by the National Natural Science Foundation of China funding [81771921, 61901084], the Wellcome Trust [WT101957, 203148/Z/16/Z], and the Engineering and Physical Sciences Research Council (EPSRC) [NS/A000027/1, NS/A000049/1]. TV is supported by a Medtronic / Royal Academy of Engineering Research Chair [RCSRF1819/7/34].