Bidirectional RNN-based Few Shot Learning for 3D Medical Image Segmentation

Few shot learning aims to learn transferable knowledge that can be used to generalize to unseen tasks with scarce labeled training data. It has been widely used for classification (Snell, Swersky, and Zemel 2017; Sung et al. 2018; Garcia and Bruna 2017; Liu et al. 2018), segmentation (Li et al. 2020; Wang et al. 2019; Zhang et al. 2019; Shaban et al. 2017; Rakelly et al. 2018), and regression tasks (Finn, Abbeel, and Levine 2017; Wang, Ramanan, and Hebert 2017; Zhou et al. 2019). For the segmentation task, (Shaban et al. 2017) first proposed a method with a fully convolutional network (FCN) that learns optimal parameters from a support image and its label to segment the relevant object from in a query image. (Dong and Xing 2018) defined a prototype, i.e., a feature vector with high-level discriminative information, through encoder and global average pooling from support data, then predicted segmentation based on similarity with the prototype of the query image. (Wang et al. 2019) additionally proposed an alignment loss that inversely guesses the support label by using the query image and the predicted query label as support data in order to select a more accurate prototype. (Zhang et al. 2019) proposed a model that refines the segmentation result by optimizing the model in an iterative manner, and (Liu et al. 2020) improved the performance by defining the prototype as a set of part-prototypes so that parts of an object can be recognized. However, prototype-based methods are limited in that the resolution of the predicted segmentation is often low since the relationship between the prototypes or parameters for prediction is learned in the down-sampled embedding space and then rapidly up-sampled using interpolation.

Recently, (Li et al. 2020) progressively performed upsampling of encoded features using a decoder with skip connections at different levels. In (Hu et al. 2019), segmentation was performed using features obtained at various stages of the encoder with an attention mechanism. Also, (Roy et al. 2020) proposed a few shot method for 3D CT organ segmentation using dense connections with squeeze and excitation blocks added between the modules for support and query data. However, most of the proposed methods have limitations in obtaining smooth segmentation for 3D organs because they estimate segmentation of the query image by only relying on support data without considering contextual information between adjacent slices. Unlike the existing methods, we propose a 3D few shot segmentation method that can consider the relationship between adjacent slices for more reliable predictions.

A number of deep learning based segmentation methods have been proposed for 3D medical image analysis. Although some architectures are fully implemented with 3D convolutions, many methods predict 3D segmentation based on 2D slices (Zhao, Chen, and Cao 2019; Oktay et al. 2018; Roth et al. 2015) or 3D patches (Korez et al. 2016; Zhu et al. 2017; Luna and Park 2018) and then perform aggregation due to the expensive computations of 3D operations and limited training data. Though such methods can learn complex tasks relatively well, the main disadvantage is the lack of global context information used for prediction. To alleviate this problem, several methods have been proposed to consider consistency between adjacent slices using recurrent neural networks (RNN). In this setting, final segmentation is achieved by obtaining predictions on adjacent slices via U-net or FCN, then perform updates on the feature maps through RNNs. For example, (Xu et al. 2019) and (Cai et al. 2017) feed U-Net predicted segmentation maps to an LSTM module. On the other hand, (Chen et al. 2016), (Cai et al. 2018), and (Li et al. 2019) used a bidirectional LSTM. (Bai et al. 2018) integrated a convolution based bidirectional LSTM into U-Net in order to model the relationship of encoded features between adjacent slices. Though successful, previous methods require a large amount of labeled training data to learn a robust model. Instead, we propose to integrate an RNN in the few shot setting for segmentation. In this way, we alleviate the issue of low training data and learn robust features considering context between slices for several organ segmentation tasks.

Transfer learning has been shown to improve model performance for different tasks by leveraging deep models already trained on larger datasets with related tasks or characteristics. Recently, several methods have been proposed to address fine tuning in the few shot setting with limited target support samples. For example, (Caelles et al. 2017) proposed a fine tuning approach for 1-shot video object segmentation using the first frame. This method was also used in metric-based few shot segmentation (Shaban et al. 2017) even though it was limited in application i.e., used to update only part of the model. In addition, optimization based few shot methods also employ fine-tuning in the intermediate stages of training. (Finn, Abbeel, and Levine 2017) proposed to temporarily update the model using support data and minimize the loss for each task, in turn good performance is obtained with only a few model updates. (Sun et al. 2019) proposed to only update the fine tuning module during the fine-tuning stage by disentangling learning of general and transferable knowledge in the modules. Inspired by prior methods, we introduce a method to learn optimal parameters for the target task by using a fine tuning process that performs additional updates by randomly dividing support data in the $K$-shot setting.

A few shot segmentation model $FSS_{\theta}$ learns parameters ${\theta}$ to segment a target object in a query image $I_{q}$ using $K$ pairs of support images and labels $\begin{Bmatrix}{I^{1}_{s},L^{1}_{s}}\end{Bmatrix}$, $\begin{Bmatrix}{I^{2}_{s},L^{2}_{s}}\end{Bmatrix}$, …, $\begin{Bmatrix}{I^{K}_{s},L^{K}_{s}}\end{Bmatrix}$, where $K$ defines the degree of supervision. Common architectures for 2D image based few shot segmentation are shown in Fig. 1 (a) to (c). (a) shows a prototypical network which defines prototypes for target object and background, and then performs segmentation using the distance to the defined prototypes, (b) is a relation network consisting of encoders and decoders for segmentation in a fully convolutional network, and (c) represents a network with dense connections between modules for support and query data. All of these methods learn the relations between support and query data of various segmentation tasks and utilize them to predict a target label in $I_{q}$ using $K$ support samples as:

(Roy et al. 2020) extended this concept to 3D image few shot segmentation. They divided a query volume $\textbf{I}_{q}$ into multiple 2D slices $I_{q,1}$, $I_{q,2}$, …, $I_{q,T}$ where $T$ is the number of slices in the axial view, and then segmented each query slice $I_{q,t}$ separately using the corresponding support slice $I_{s,u_{t}}$ and label $L_{s,u_{t}}$, where $u_{t}$ is the index of the support sample. To determine $I_{s,u_{t}}$ and $L_{s,u_{t}}$ in a 3D support volume $\textbf{I}_{s}$ with label $\textbf{L}_{s}$, it was assumed that the starting and ending slice locations of the organ of interest in $\textbf{I}_{q}$ and $\textbf{I}_{s}$ were known. The index of a support slice corresponding to $I_{q,t}$ is obtained via $u_{t}=round((t/T)\times\hat{T})$, where $\hat{T}$ is the number of interest slices in $\textbf{I}_{s}$. This assumption is reasonable since the organ of interest is in a similar position from person to person, e.g., the liver is always located in the upper right part of the abdomen, even though it varies in size and shape. In this setting, few shot segmentation of a 3D image can be represented by the following formulation:

Most 2D based few shot segmentation models can follow this setting, but the relation between adjacent slices is not considered. In this study, we follow the mentioned problem setting but propose to incorporate adjacent slice information to accurately segment $I_{q,t}$. Formally,

where $2n_{a}+1$ is the number of adjacent slices and $t_{0}$ is the index of the center of multiple slices. To achieve this, we introduce a 3D few shot segmentation method using the bidirectional RNN.

Our model performs segmentation in three stages; $(1)$ features of support and query images are extracted through two separate encoders $E_{s}$ and $E_{q}$, respectively. $(2)$ A bidirectional GRU models the relationship between features extracted from adjacent slices. $(3)$ Finally, using updated feature maps and low level features in $E_{q}$, the decoder predicts the segmentation. An overview of the one-shot segmentation model is shown in Fig. 1 (d).

Since the target organ is never seen in the training stage, it may be challenging for model to learn the optimal parameters for the target. Thus, in multi-shot scenarios, we perform transfer learning with few target-support samples prior to testing. Specifically, we temporarily sample support and query pairs from the support data and update our model. For example, in $K$ shot testing stage, we collect existing pairs of support and query data by choosing $K-1$ samples from the $K$ support set as temporary support samples and use the remainder as a temporary query image to retrain our model. Since we use $2n_{a}+1$ adjacent slices from a 3D volume, many different training pairs can be sampled enabling a robust fine tuning process. Moreover, to avoid overfitting and encourage training stability, we use random flipping and rotation based augmentation during training. Using this strategy, our model can effectively adapt to the novel characteristics of target data.

The proposed method was evaluated on the Multi-atlas labeling Beyond the Cranial Vault (BCV) dataset (Landman et al. 2017) which was used in a challenge at MICCAI 2015. The dataset includes 30 3D CT scans with segmentation labels for 15 organs. Among 15 organs, the labels of 9 organs (spleen, left kidney, esophagus, liver, stomach, aorta, inferior vana cava, bladder, and uturus) were used in our experiments since the other 6 organs were too small or hard to segment even with a supervised learning method due to large shape variations.

Moreover, we used two external datasets to see if the proposed model would be applicable to data with different characteristics. We employed CTORG (Blaine Rister and Rubin 2019) which contains 119 images with labels of 6 organs (lung, bones, liver, left and right kidneys, and bladder). It is worth noting that variations of in-plane resolution and thickness between images were significant since the dataset was collected from multiple sites. In our experiments, external tests were performed on the liver, kidney, and bladder, excluding the lung and brain since many CT scans did not include the whole part of the lung, and the brain samples were limited. Second, we also evaluated our method on the DECATHLON(Simpson et al. 2019) dataset. It consists of several images with 10 different organs (CT Liver, multimodal MRI brain tumors, mono-modal MRI hippocampus, CT lung tumors, Multimodal prostate, monomodal left atrium, CT pancreas, CT colon cancer primaries, CT hepatic vessels, and CT spleen). In a similar fashion with earlier experiments, spleen and liver data were used in the external tests, excluding organs that were too small or had severe shape changes.

The BCV dataset was divided into 15 volumes for training or selecting support data, 5 volumes for validation, and 10 volumes for testing for each organ. In training, pairs of support and query data were randomly sampled from the 15 volumes with 8 organs except a certain target organ to train the few shot models. For testing, support data was randomly sampled among the 15 volumes for the target organ, while the 10 volumes were used as query images. Since performing experiments for all organs was time-consuming, we tested our model on 4 clinically important organs (spleen, liver, kidney, and bladder) that are not too small. For example, the adrenal gland was excluded because it appears in limited slices of CT scans and was often hard to figure out the organ’s 3D structure. For external validation, the models trained on the BCV dataset were applied to 65 liver, 63 kidney, and 53 bladder samples in CTORG, and 27 spleen and 87 liver data in the DECATHLON dataset. Voxel intensities of all images were normalized to range from 0 to 1 and slices were cropped into squares with the same size for each organ, then resized to 256$\times$256.

To show the effectiveness of proposed model, we compared our method with a U-net based supervised method and three few shot models ($FSS_{base}$ (Li et al. 2020), $FSS_{prototype}$ (Wang et al. 2019), and $FSS_{SE}$ (Roy et al. 2020)) recently proposed. A U-net (Ronneberger, Fischer, and Brox 2015) trained with only one sample per organ was used as a lowerbound while that trained with all accessible data was used as a upperbound model. For fair comparison with our proposed model, U-net was modified to use 5 adjacent axial slices as input and consisted of 2D convolutional encoder and decoder. He initialization (He et al. 2015) was used for all models with Adam (Kingma and Ba 2014) optimization and a learning rate of $10^{-4}$. For every iteration in the training stage, support and query volumes were randomly selected from training data containing various organ segmentation labels except the target organ. For the bidirectional GRU models, a total 5 slices were feed into the model, i.e., $n_{a}$ was set as 2. Similar settings were used for inference. In addition, the same parameter initialization and data augmentation (flipping and rotation) was applied across all evaluated models.

$FSS_{base}$ is a baseline model with a similar architecture to the proposed model if the bidirectional GRU module is omitted. $FSS_{prototype}$ uses prototypes and an alignment process for predictions. It defines prototypes of foreground and background to implement distance-based pixel-wise classification on reduced feature maps extracted by an encoder. On the other hand, $FSS_{SE}$ model uses squeeze and excitation blocks with skip connections trained from scratch, with a separate encoder and decoder for support and query data. Except for $FSS_{SE}$, we evaluated 1, 3, and 5 shot models on the internal and external testing datasets. Since $FSS_{SE}$ was designed for the one-shot setting, 3 and 5 shot settings were not considered for evaluation. As for the proposed model, we note it as $FSS_{BiGRU}$.

Since the performance of segmentation may vary depending on how the support set is selected, the experiment was performed with different support sets randomly sampled 5 times for each query sample, and we report the average across trials. Segmentation performance was measured by the dice similarity score between the prediction and label.