An application of cascaded 3D fully convolutional networks for medical image segmentation

Introduction

Recent advances in fully convolutional networks (FCN) have made it feasible to train models for pixel-wise segmentation in an end-to-end fashion (Long et al., 2015). Efficient implementations of 3D convolution and growing GPU memory have made it possible to extent these methods to 3D medical imaging and train networks on large amounts of annotated volumes. One such example is the recently proposed 3D U-Net (Çiçek et al., 2016), which applies a 3D FCN with skip connections to sparsely annotated biomedical images. Alternative architectures for processing volumetric images have also been successfully applied to 3D medical image segmentation (Milletari et al., 2016; Chen et al., 2016; Dou et al., 2017). In this work, we show that a 3D FCN, like 3D U-Net, trained on manually labeled data of several anatomical structures (ranging from the large organs to thin vessels) can also achieve competitive segmentation results on clinical CT images, very different from the original application of 3D U-Net using confocal microscopy images. We furthermore compare our approach to 2D FCNs applied to the same images.

Our approach applies 3D FCN architectures to problems of multi-organ and vessel segmentation in a cascaded fashion. A FCN can be trained on whole 3D CT scans. However, because of the high imbalance between background and foreground voxels (organs, vessels, etc.) the network will concentrate on differentiating the foreground from the background voxels in order to minimize the loss function used for training. While this enables the FCN to roughly segment the organs, it causes particularly smaller organs (like the pancreas or gallbladder) and vessels to suffer from inaccuracies around their boundaries.

To overcome this limitation, we learn a second-stage FCN in a cascaded manner that focuses more on the boundary regions. This is a coarse-to-fine approach in which the first-stage FCN sees around 40% of the voxels using only a simple automatically generated mask of the patient’s body. In the second stage, the amount of the image’s voxels is further reduced to around 10%. In effect, this step narrows down and simplifies the search space for the FCN to decide which voxels belong to the background or any of the foreground classes; this strategy has been successful in many computer vision problems (Viola and Jones, 2004; Li et al., 2016). Our approach is illustrated on a training example in Fig. 1.

Multi-organ segmentation has attracted considerable interest over the years. Classical approaches include statistical shape models (Cerrolaza et al., 2015; Okada et al., 2015), and/or employ techniques based on image registration. So called multi-atlas label fusion (Rohlfing et al., 2004; Wang et al., 2013; Iglesias and Sabuncu, 2015) has found wide application in clinical research and practice. Approaches that combine techniques from multi-atlas registration and machine learning are also common place and have been successfully applied to multi-organ segmentation in abdominal imaging (Tong et al., 2015; Oda et al., 2016). However, a fundamental disadvantage of image registration based methods is there extensive computational cost (Iglesias and Sabuncu, 2015). Typical methods need hours of computation time in order to complete on single desktop machines (Wolz et al., 2013).

The recent success of deep learning based classification and segmentation methods are now transitioning to applications of multi-class segmentation in medical imaging. Recent examples of deep learning applied to organ segmentation include (Roth et al., 2017; Zhou et al., 2016b; Christ et al., 2016; Zhou et al., 2016a). Many methods focus on the segmentation of single organs like prostate (Milletari et al., 2016), liver (Christ et al., 2016), or pancreas (Roth et al., 2015, 2016b). Multi-organ segmentation in abdominal CT has also been approached by works like (Hu et al., 2017; Gibson et al., 2017). Most methods are based on variants of FCNs (Long et al., 2015) that either employ 2D convolutional layers in a slice-by-slice fashion Roth et al. (2016b); Zhou et al. (2016b); Christ et al. (2016); Zhou et al. (2016a), 2D convolutions on orthogonal (2.5D) cross-sections (Roth et al., 2015; Prasoon et al., 2013), and 3D convolutional layers (Milletari et al., 2016; Chen et al., 2016; Dou et al., 2017; Kamnitsas et al., 2017). A common feature of these novel segmentation methods is that they are able to extract the features useful for image segmentation directly from the training imaging data, which is crucial for the success of deep learning (LeCun et al., 2015). This avoids the need for hand-crafting features that are suitable for detection of individual organs.

2 Contributions

Due to the automatic learning of image feature and in contrast to previous approaches of multi-organ segmentation where separate models have to be created for each organ (Oda et al., 2016; Tong et al., 2015), our proposed method allows us to use the same model to segment very different anatomical structures such as large abdominal organs (liver, spleen), but also vessels like arteries and veins. Furthermore, other recent FCN-based methods that applied in medical imaging in cascaded/iterative fashion were often constrained to using rectangular bounding boxes around single organs (Roth et al., 2017; Zhou et al., 2016b) and/or performing slice-wise processing in 2D (Christ et al., 2016; Zhou et al., 2016a).

Methods

Convolutional neural networks have the ability to solve challenging classification tasks in a data-driven manner. Given a training set of images and labels $\mathbf{S}=\left\{(I_{n},L_{n}),n=1,\dots,N\right\}$, $I_{n}$ denotes the raw CT images and $L_{n}$ denotes the ground truth label images. Each $L_{n}$ contains $K$ class labels consisting of the manual segmentations of the foreground anatomy (e.g. artery, portal vein, lungs, liver, spleen, stomach, gallbladder, and pancreas) and the background for each voxel in the CT image. Our employed network architecture is the 3D extension by Çiçek et al. (2016) of the U-Net proposed by Ronneberger et al. (2015). U-Net, which is a type of fully convolutional network (FCN) (Long et al., 2015) was originally proposed for bio-medical image applications, utilizes deconvolution (Long et al., 2015) (or sometimes called up-convolutions (Çiçek et al., 2016)) to remap the lower resolution feature maps within the network to the denser space of the input images. This operation allows for denser voxel-to-voxel predictions in contrast to previously proposed sliding-window CNN methods where each voxel under the window is classified independently making such architecture inefficient for processing large 3D volumes. In 3D U-Net, operations such as 2D convolution, 2D max-pooling, and 2D deconvolution are replaced by their 3D counterparts (Çiçek et al., 2016). We use the open-source implementation of 3D U-Nethttp://lmb.informatik.uni-freiburg.de/resources/opensource/unet.en.html based on the Caffe deep learning library (Jia et al., 2014). The 3D U-Net architecture consists of analysis and synthesis paths with four resolution levels each. Each resolution level in the analysis path contains two $3\times 3\times 3$ convolutional layers, each followed by rectified linear units (ReLU) and a $2\times 2\times 2$ max pooling with strides of two in each dimension. In the synthesis path, the convolutional layers are replaced by deconvolutions of $2\times 2\times 2$ with strides of two in each dimension. These are followed by two $3\times 3\times 3$ convolutions, each of which has a ReLU. Furthermore, 3D U-Net employs shortcut (or skip) connections from layers of equal resolution in the analysis path to provide higher-resolution features to the synthesis path (Çiçek et al., 2016). The last layer contains a $1\times 1\times 1$ convolution that reduces the number of output channels to the number of class labels $K$. This architecture has over 19 million learnable parameters and can be trained to minimize a weighted voxel-wise cross-entropy loss (Çiçek et al., 2016). A schematic illustration of 3D U-Net is shown in Fig. 2.

The voxel-wise cross-entropy loss is defined as

where $\hat{p}_{k}$ are the softmax class probabilities

$N$ are the total number of voxels $x$, $S_{k}$ is the set of voxels within one class in $L_{n}$, and $k\in[1,2,\dots,K]$ indicates the ground truth class label. The input to this loss function is real valued output predictions $x\in[-\infty,+\infty]$ from the last convolutional layer.

However, in most cases minimizing this loss will instantly make the network converge to classifying every voxel as background. This is because of the large dominance of the background class in the images. In order to combat this large data imbalance between foreground/background voxels and differently sized organs and vessels, we apply a voxel-wise weight $\lambda_{k}$ to this loss function (Eq. 1). In this work, we choose $\lambda_{i}$ such that $\sum^{K}_{k=1}\lambda_{k}=1$, with

where $N_{k}$ is the number of voxels in each class $S_{k}$, and $N_{C}$ is the number of voxels within a candidate region $C_{1}$ or $C_{2}$. The weights $\lambda_{i}$ help to balance the common voxels (i.e., background) with respect to such smaller organs as vessels or the pancreas by giving more weight to the latter.

Now, the weighted cross-entropy loss can be written as:

We use the loss formulation in Eq. 4 for all experiments in this paper.

2 Coarse-to-fine prediction

In the first stage, we apply simple thresholding in combination with morphological operations (hole filling and largest component selection) to get a mask of the patient’s body. This mask can be utilized as candidate region $C_{1}$ to reduce the number of voxels necessary to compute the network’s loss function and reduce the amount of input 3D regions shown to the CNN during training to about 40%.

nd Stage

After training the first-stage FCN, it is applied to each image to generate candidate regions $C_{2}$ for training the second-stage FCN (see Fig. 1). We define the predicted organ labels in the testing phase using the $\mathop{\rm argmax}\limits$ of the class probability maps. All foreground labels are then dilated in 3D using a voxel radius of $r$ in order to compute $C_{2}$, resulting in a binary candidate map.

When comparing the recall and false-positive rates of the first-stage FCN with respect to $r$ for both the training and validation sets, $r=3$ gives good trade-off between high recall ($>$99%) and low false-positive rates ($\sim$10%) for each organ on our training and validation sets (see Fig. 6).

Our overall multi-stage training scheme is illustrated in Fig. 3

3 Training

The network iteratively adjusts its parameters by stochastic gradient descent. Batch normalization is used throughout the network for improved convergence and we utilize random elastic deformations in 3D during training to artificially increase the amount of available data samples and increase robustness, similar to (Çiçek et al., 2016). Hence, we randomly sample deformation fields from a uniform distribution with a maximum displacement of $\pm$4 and a grid spacing of 32 voxels (see Fig. 4). Furthermore, we applied random rotations between $-5^{\circ}$ and $+5^{\circ}$, and translations of -20 to 20 voxels in each direction at each iteration in order to generate plausible deformations during training. Each training sub-volume is randomly extracted from $C_{1}$ or $C_{2}$ in both stages.

4 Testing

The CT image is processed by the 3D FCN using a tiling strategy (sliding-window) (Çiçek et al., 2016) as illustrated in Fig. 5. For greater speed, we use non-overlapping tiles in the first stage and investigate the use of non-overlapping and overlapping tiles in the second. When using overlapping tiles (with a $4\times$ higher sampling rate of each voxel $x$), the resulting probabilities for the overlapping voxels are averaged:

Experiments & Results

Our dataset includes 331 contrast-enhanced abdominal clinical CT images in the portal venous phase used for pre-operative planning in gastric surgery. Each CT volume consists of $460-1177$ slices of $512\times 512$ pixels. The voxel dimensions are [0.59-0.98, 0.59-0.98, 0.5-1.0] mm. A random split of 281/50 patients is used for training and validating the network, i.e., determining when to stop training to avoid overfitting. In order to generate plausible deformations during training, we sample from a normal distribution with a standard derivation of 4 and a grid spacing of 32 voxels, and apply random rotations between $-5^{\circ}$ and $+5^{\circ}$ to the training images. No deformations were applied during testing. We trained 200,000 iterations in the first stage and 115,000 in the second. Table 1 summarizes the Dice similarity scores for each organ labeled in the 50 validation cases. On average, we achieved a 7.5% improvement in Dice scores per organ. Small, thin organs such as arteries especially benefit from our two-stage cascaded approach. For example, the mean Dice score for arteries improved from 59.0 to 79.6% and from 54.8 to 63.1% for the pancreas. The effect is less pronounced for large organs, like the liver, the spleen, and the stomach. Fig. 7 shows an example result from the validation set and illustrates the tiling approach. The 3D U-Net separates the foreground organs well from the background tissue of the images.

2 Testing

Our test set is different from our training and validation data. It originates from a different hospital, scanners, and research study with gastric cancer patients. 150 abdominal CT scans were acquired in the portal venous phase. Each CT volume consists of $263-1061$ slices of $512\times 512$ pixels. Voxel dimensions are [0.55-0.82, 0.55-0.82, 0.4-0.80] mm. The pancreas, liver, and spleen were semi-automatically delineated by three trained researchers and confirmed by a clinician. Figure 8 shows surface renderings for comparison of the different stages of the algorithm. A typical testing case in the first and second stages is shown using non-overlapping and overlapping tiles. Dice similarity scores are listed in Table 2. The second stage achieves the highest reported average score for pancreas in this dataset with 82.2% $\pm$10.2%. Previous state of the art on this dataset was at 75.1% $\pm$15.4% while using leave-one-out-validation (Oda et al., 2016).

The testing dataset provides slightly higher image quality than our training/validation dataset. Furthermore, its field of view is more constrained to the upper abdomen. This likely explains the improved performance for liver and pancreas compared to the validation set in Table 1.

3 Comparison to other methods

Even though direct comparison is difficult due to the differences in datasets, training/testing evaluation schemes, and segmented organs, we try to indicate how well our model performed with respect to recent state-of-the-art methods in Table 3. In particualar, we provide a comparison to recent methods on two different datasets: (1) our own in-house dataset for pancreas segmentation, acquired at Nagoya University Hospital, Japan, and consisting of 150 CT images; and (2) the publicly available TCIA Pancreas-CT dataset of 82 patient imageshttps://wiki.cancerimagingarchive.net/display/Public/Pancreas-CT (Roth et al., 2016a) hosted by TCIA (Clark et al., 2013). (Roth et al., 2016a). For comparison with (2), we use the same 4-fold cross-validation (CV) split as in (Roth et al., 2015, 2017).

Our results on dataset (1) achieves the highest reported performance in testing. On the other hand, our results on the public dataset (2) are comparable to other recent works that developed methods especially targeting this dataset and focusing on pancreas segmentation alone (Roth et al., 2017; Zhou et al., 2016b).

4 Direct comparison to 2D FCN networks

Furthermore, we implement the method of Zhou et al. (Zhou et al., 2016a, 2017) and apply it to the same dataset. This method employs a combination of three 2D FCNs trained on the orthogonal planes of the images. The results of each model are then fused by majority voting. This dataset consists of 240 3D CT scans with 18 manually annotated organs. A split of 228/12 cases was used for our training/testing as in (Zhou et al., 2016a, 2017). A direct comparison can be seen in Table 4. It can be observed that our 3D FCN approach has a clear advantage for the smaller, thinner organs (like aorta, esophagus, gallbladder, inferior vena cava, portal vein, and prostate) but only performs comparable to the 2D FCNs when aiming at the larger organs (like lungs, liver, kidneys). Furthermore, a slightly higher overall performance can be observed for the average of all organ/vessel predictions when using the proposed cascaded 3D FCN approach.

5 Computation

Fine-tuning to other datasets

One advantage of deep learning based models is their ability to transfer learned features across dataset domains (Shin et al., 2016). To this end, we trained a general FCN model employing the 3D U-Net architecture (Çiçek et al., 2016) on the large dataset of CT scans including the major abdominal organ labels of Section 3.1. This model can then be fine-tuned to other (smaller) datasets aiming at more detailed classification tasks or different field of views. For this purpose, we utilize separate training, fine-tuning, and testing datasets. As mentioned above, the general training set consists of 280 clinical CT images with seven abdominal structures (artery, vein, liver, spleen, stomach, gallbladder, and pancreas) labeled.

We then fine-tune on a much smaller dataset consisting only of 20 contrast enhanced CT images from the Visceral Challenge datasethttp://www.visceral.eu/benchmarks/anatomy3-open/ (Jimenez-del Toro et al., 2016), but with substantially more anatomical structures labeled in each image (20 in total). This fine-tuning process across different datasets is illustrated in Fig. 9 with some ground truth label examples used for pre-training and fine-tuning. In fine-tuning, we use a 10 times smaller learning rate. We furthermore test our models on a completely unseen data collection of 10 torso CT images with 8 labels, including organs that were not labeled in the original abdominal dataset, e.g. the kidneys and lungs. A probabilistic output for kidney (not in the pre-training dataset) from our model is shown in Fig. 10.

In testing, we deploy our fine-tuned model using a non-overlapping tiling approach as in previous sections. An automated segmentation result on the unseen test dataset by our fine-tuned model is shown in Fig. 11. Our fine-tuned approach provides a Dice score of right lung, left lung, liver, gall bladder, spleen, right kidney, left kidney, and pancreas are 0.96, 0.97, 0.95, 0.77, 0.90, 0.90, 0.88, and 0.36, respectively (summarized in Table 5). The relatively lower score for pancreas is due to several outlier cases on this dataset. These outliers are likely caused by variations of contrast enhancement across the datasets and the higher variability of the pancreas’ shape and intensity profile compared to other organs across different patients.

Our approach and results, however, illustrate the generalizability and robustness of our models across different datasets. Fine-tuning can be useful when the amount of training examples for some target organs are limited. In this case, transfer learning achieves slight improvements over learning from scratch, especially in the kidneys (see Table 5). It should be noted that for this particular application, data augmentation already gives a good performance when learning models from scratch.

Discussion

The cascaded coarse-to-fine approach presented in this paper provides a simple yet effective method for employing 3D FCNs in medical imaging settings. No post-processing was applied to any of the FCN outputs. The improved performance stemming from our cascaded approach is especially visible in smaller, thinner organs, such as arteries and veins, particularly when compared to other recent FCN approaches using 2D FCNs (Zhou et al., 2016a). Our results and recent literature indicate that 2D FCNs and especially the combination of orthogonally applied 2D FCNs (Zhou et al., 2016a; Roth et al., 2017; Zhou et al., 2016b) might be sufficient for larger and mid-sized organs. In fact, the combination of 2D FCNs even slightly outperforms our 3D approach for some organs. On the other hand 3D convolutional kernels are important for distinguishing the thin (vessel-like) and small organs as can be seen in the improved performance of our approach. When compared to other cascaded approaches using 2D FCNs that focus on single organs (Roth et al., 2017; Zhou et al., 2016b), we perform similar to the state of the art. Our findings are also consistent with (Roth et al., 2017; Zhou et al., 2016b) that show that cascaded approaches are useful for applying deep learning methods to medical image segmentation. Note that we used different datasets (from different hospitals and scanners) for separate training/validation and testing. These experiments illustrate our method’s generalizability and robustness to differences in image quality and populations. Running the algorithms at half resolution allows efficient training on a single GPU. In contrast, using the same field of view for each subvolume with the original resolution would require $8\times$ more memory with the current architecture and would force us to reduce the amount of context visible to the 3D FCNs. In this work, we utilized 3D U-Net for the segmentation of CT scans. However, the proposed cascaded approach in principle should also work well for other 3D CNN/FCN architectures and 3D image modalities. Exploration of other loss functions such as the Dice score (Milletari et al., 2016; Li et al., 2017) could help further in dealing with the class imbalance issue. We used Caffe’s stochastic gradient descent solver (Jia et al., 2014) for all experiments in this work. Alternative optimizers could further improve training performance (Kingma and Ba, 2014).

In the future, prediction results from different models could be combined in order to achieve the best overall performance. Furthermore, additional anatomical constraints could be included in order to guarantee topologically correct segmentation results (BenTaieb and Hamarneh, 2016; Oktay et al., 2017). With growing amounts of available GPU memory, the need for computing overlapping sub-volume predictions as in this work will be reduced as it will be come possible to reshape the network to accept arbitrary 3D input image sizes (Long et al., 2015).

Conclusion

In conclusion, we showed that a cascaded deployment of volumetric fully convolutional networks (3D U-Net) can produce competitive results for medical image segmentation on a clinical CT dataset while being efficiently deployed on a single GPU. An overlapping tiles approach during testing produces better results with only moderate additional computational cost. The proposed method compares favorably to recent state-of-the-art work on a completely unseen dataset. Our results indicate that 3D convolutional features are advantageous for detecting smaller organs and vessel. A promising future direction might be hybrid approaches that combine 2D and 3D FCN-type architectures at multiple scales. We have made our code, pre-trained models, and fine-tuned models available for downloadhttps://github.com/holgerroth/3Dunet_abdomen_cascade in order to allow further applications and fine-tuning to different datasets. Acknowledgments This paper was supported by MEXT KAKENHI (26108006, 26560255, 25242047, 17H00867, 15H01116) and the JPSP International Bilateral Collaboration Grant.

Conflict of interest statement: The authors declare that they have no conflict of interest.

References

References