FocusNetv2: Imbalanced Large and Small Organ Segmentation with Adversarial Shape Constraint for Head and Neck CT Images

Recently, convolutional neural networks have greatly advanced the field of medical image analysis because of its ability to learn more representative features from data. CNNs demonstrate state-of-the-art performance in many challenging tasks, such as image classification, segmentation, detection, registration, super-resolution and etc.

Long et al. [25] first proposed fully convolutional network (FCN), which uses convolutions with 1$\times$1 sized filter to replace the fully connected layer, and allows the prediction of multiple pixels at the same time. Ronneberger et al. [37] further built a “U” shaped network (named U-Net) with a contracting path and a symmetric expanding path. Skip connections are also used to propagate features from early layers to later ones. A vast number of works based on variants of FCN and U-Net are applied in the field of 2D medical image segmentation [5, 52, 3, 38, 40].

For 3D images like CT or MR, 2D CNNs can be utilized in a slice-by-slice manner, however, the context information encoded in the volumetric data is ignored. Some 2.5D methods [39, 48] attempted to incorporate 3D spatial information by using three orthogonal slices or adjacent slices. But their representation capabilities are still limited by 2D convolutional kernels. To overcome this weakness, 3D CNN-based algorithms are proposed. For example, the 3D version of U-Net is proposed by Çiçek et al. [6]; Milletari et al. [27] proposed V-Net, which introduces the residual connection [17] between building blocks to mitigate gradient vanishing problem. Several 3D networks were also proposed for different applications, such as Merkow et al. [26], Dou et al. [10], Kamnitsas et al. [20].

Even though 3D CNN based methods can better leverage spatial context for learning better feature representations, the sample imbalanced problem is magnified in 3D tasks, as the training errors are mostly dominated by voxels belong to large organs. Ronneberger et al. [37] proposed to use the weighted cross-entropy loss function, while Milletari et al. [27] proposed the Dice coefficient loss, they can only mitigate the challenge of unbalanced data but are far away from solving it.

There are numerous works of OAR segmentation proposed for radiotherapy treatment planning of different body parts. Atlas-based methods are among the most commonly used traditional approaches. The optimal transformation between the atlas, which has the pre-segmented annotation map, and the image to be segmented is aligned by affine and deformable registration. Then the segmentation for the target image can be obtained by applying this transformation on the annotation map of the reference image. The reference images can be multiple ones with expert annotations or templates generated from the training set. The accuracy of atlas-based approaches is affected by two factors: First, the capability of the registration method, whether it can align the target image and atlas images accurately. Different approaches, such as Demons registration [42, 33], block-matching [32, 15], and B-Spline registration [53], have been proposed. Second, physiologic or pathologic anatomical variations of some organs make it difficult to find the optimal correspondences between target images and reference images, thus some methods are proposed to use atlases that reflect average patient anatomy [8] or the fusion of results from multi-atlas [7, 36]. Some hybrid methods post-process results from atlas-based segmentation using active contours [53] and graph cuts [23, 11]. Although atlas-based methods have the advantages of robustness and can perform segmentation without user interaction, they are based on image registration techniques and might generate incorrect organ maps if the organs are occupied by tumors. The time cost can be up to tens of minutes due to its huge amount of computation.

Recently, convolutional neural networks were adopted to advance OAR delineation accuracy significantly . Ibragimov and Xing [19] proposed the first deep learning-based algorithm. They first detect OARs by the head center point and then train a patch-based CNN to classify voxels in the region of interest. Ren et al. [35] proposed an interleaved 3D-CNN for the joint segmentation of small organs in HaN, where the region of interest is obtained via registration techniques. Zhu et al. [55] proposed a 3D squeeze-and-excitation U-Net for fast segmentation. Tong et al. [43] proposed a fully convolutional neural network with a shape representation model. Tang et al. [41] also proposed a two-stage segmentation method based on detection, where local contrast normalization is applied in the segmentation head for achieving better segmentation performance. There are also methods [28] segment OARs in other modalities, such as MRI.

Current CNN-based approaches are usually trained with pixel-wise losses, such as the cross-entropy loss and the dice loss, which consider pixels’ prediction errors separately. Even with large enough receptive fields, they cannot maintain the overall shapes or higher-level structures of the organ of interest. Conditional random field (CRF) and graph cut methods were proposed to enforce spatial contiguity of the output label maps. Several approaches attempted to incorporate shape constraints in CNNs, Mosinska et al. [29] incorporated a pre-trained VGG network to extract higher-order topological features from the network’s predictions and ground truth label maps, and then minimizes the $L2$ loss between them. However, the ImageNet pre-trained VGG network may not capture various organs’ anatomy in medical images. Oktay et al. [31] used an autoencoder (AE) to learn representations of shapes from ground-truth annotations and minimize the Euclidean loss between the encoded latent codes of the network’s predicted label maps and ground truth label maps. Tong et al. [43] applied similar ideas with Oktay et al. [31] in the segmentation of OARs. Al Arif et al. [1] modified the UNet to generate a signed distance function (SDF) instead of segmentation maps and computes the errors directly in the shape domain using principal component analysis (PCA). Adversarial losses were also introduced as high-order shape regularization [50, 49]. The basic idea of methods mentioned above are similar, i.e., instead of measuring the shape similarity with ground-truth in the pixel domain, they proposed to measure the similarity in a lower dimension manifold by different kinds of projection (i.e., pre-trained VGG network, autoencoder, PCA or discriminator). Therefore, the key is the quality of the shape projection. However, existing adversarial discriminators lack generalization capability. The autoencoder in Oktay et al. [31], Tong et al. [43] was trained with only ground-truth annotation. It cannot well encode various predicted shapes by the segmentation models. To mitigate the problem, we proposed a novel adversarial shape constraint based on autoencoder that incorporates benefits from both autoencoder and adversarial discriminators.

There are three main differences between this work and the previous work. First, we designed a high-performance backbone network with fewer down-samplings but using DenseASPP [51] to retain more detailed high-resolution information as well as learn multi-scale features. Second, we use a two-stage segmentation model for small organs to solve the imbalance problem between large and small organs. The third and most crucial difference is that we use the adversarial autoencoder to apply shape constraints to the network prediction, which allows our network to generate predicted shapes that conform to prior medical knowledge for specific organs with unclear boundaries and low contrast. In practical clinical applications, our model can generate delineation results that are more acceptable by human doctors to improve the efficiency of radiotherapy treatment planning.

In our task, the ratio between the largest and the smallest organ can reach nearly ${500:1}$, which makes the loss dominated by the large number of large-organ voxel samples. In recent literature, focal loss [24] and generalized dice loss are two effective loss functions that try to solve the problem of class imbalance. We propose to use a weighted focal loss for multi-class segmentation,

where $C$ is the number of categories, $p_{t}$ is the probability of class $t$. For training the S-Net, $\alpha_{t}$ is the weight of each organ, which is inversely proportional to each organ’s average size. ${(1-p_{t})^{\gamma}}$ is the modulating factor that weights less on easy samples (voxels) with prediction confidence $p_{t}$ close to $1$ and weight larger on challenging samples. The focal loss can adaptively down-weight the loss contributed from easy samples, while suppress lightly the contributions of incorrectly classified hard samples. In our experiment, $\gamma$ is empirically set as $2$.

Generalized dice loss is another loss function that directly optimizes towards the evaluation metrics. We adopt the following generalized dice loss,

where $y_{t}$ and $p_{t}$ denote the ground-truth labels and predicted probabilities of class $t$. In our experiment, the combination of focal loss and dice loss results in the best segmentation accuracy. The total segmentation loss is therefore defined as

where $\lambda$ weights the two losses and we empirically set $\lambda=1$ in our experiments.

Due to CT’s imaging principle, some OARs do not show obvious boundaries in the images, such as the optic chiasm. As OARs are normal organs, they usually have relatively consistent shapes among different patients, incorporating high-order shape constraint to the segmentation network can make the prediction more consistent with prior anatomical knowledge. We propose a novel adversarial autoencoder (AAE) to introduce shape regularization into the training of our segmentation framework. To the best of our knowledge, this is the first segmentation method that leverages both the autoencoder and adversarial learning for constraining segmentation masks.

A good design of shape regularization term should have the following two characteristics. First, it needs to be capable of representing shapes in a differentiable way, so that the segmentation network can be trained with back-propagation for shape regularization; Second, it should be able to distinguish subtle differences between shapes, such that as the shape predicted by the segmentation network gets closer to the ground-truth shapes, it still gives a correct penalty.

The shapes represented by label maps are highly structured and high dimensional, making it extremely challenging to measure the similarity between two shapes in such high dimensional space. The high-dimensional shape usually lies in a lower-dimensional shape manifold [44, 31], where each shape would be mapped to a lower-dimensional point (vector) in the subspace. If a shape manifold is successfully discovered, starting from a point on the manifold (corresponding to one specific shape), we can traverse along with different directions on the manifold. The corresponding shape would change smoothly and continuously at the semantic level. Therefore, we measure the similarity between two shapes in a low-dimensional shape manifold. We use a shape autoencoder, which is a neural network trained to reconstruct the input organ shape as much as possible (see Fig. 4). The bottleneck structure enables the autoencoder to transform the input shape into a latent code that captures its salient features while discarding irrelevant features. Therefore, if the autoencoder can well reconstruct the input shape, the latent space is a good approximation of low-dimensional shape manifold [22, 54]. Moreover, the autoencoder is differentiable and can regularize estimated organ shapes by minimizing the distance between the predicted organ shape and ground-truth shape in the latent space.

For the second characteristic, the accurate measurement of the similarity between the predicted organ shape and the ground-truth organ shape is crucial. Therefore, we introduce an adversarial training scheme for training the adversarial autoencoder. The autoencoder is trained with both the predicted shape from the small organ segmentation branch and the corresponding ground-truth shapes. It has two loss terms, where the first one is to reconstruct the input shapes for learning shape representations by minimizing the conventional reconstruction loss,

where $x$ is the input image, $y$ is its corresponding ground-truth label, $G$ is the segmentation network, $G(x)$ is the predicted binary organ mask from SOS-Net given the input image, $D(y)$ and $D(G(x))$ are the reconstruction results of AAE, given the ground truth label $y$ and predicted organ mask $G(x)$.

The other adversarial loss term tries to distinguish the latent codes of predicted shapes and ground-truth shapes by maximizing their distance in the low-dimensional manifold. In this way, we enforce the autoencoder to better encode the two types of shapes and capture their subtle differences, while the segmentation network is encouraged to fool the autoencoder to being unable to capture the subtle differences. Therefore, the proposed adversarial shape loss is formulated as

where $D_{latent}(y)$ or $D_{latent}(G(x))$ is the latent code and multi-scale decoder features of the ground-truth organ shape $y$ and the predicted organ shape $G(x)$, as illustrated in Fig. 4. Intuitively, the segmentation network $G$ and the AAE $D$ play a min-max game. The segmentation network $G$ tries to predict masks that are consistent with shape prior to minimize the distance in low-dimensional shape manifold, while the AAE $D$ tries to learn better encoding to maximize the distance.

The overall objective function for segmentation network and with the shape regularization term from the adversarial autoencoder is therefore defined as

In practice, the segmentation network $G$ and the adversarial autoencoder $D$ are trained in an alternative optimization scheme. The $G$ is first optimized by fixing $D$ and minimizing the following loss

where $\lambda_{1}$ is empirically set as 5 to balance the two terms. Optimizing $L_{G}$ would encourage the segmentation network $G$ to output organ shapes that are consistent with the ground-truth shapes. The $D$ is then optimized when $G$ is fixed by minimizing the following loss,

where $\lambda_{2}$ is set as 0.001 in our experiment as larger $\lambda_{2}$ may cause unstable training of the proposed AAE. During the training of segmentation network $G$, the estimated organ shapes by $G$ will gradually become closer to the ground-truth organ shapes, therefore, if the encoded latent codes are not distinguishable between the estimated organ shapes and the ground-truth organ shapes, the autoencoder may hardly provide effective supervision for the segmentation network $G$. Therefore, maximizing the distance between $D_{latent}(G(x))$ and $D_{latent}(y)$ encourages the autoencoder to encode subtle difference between them. Eqs. (7) and (8) are optimized alternatively to gradually improve both the segmentation network $G$ and the adversarial autoencoder $D$.

Discussions on segmentation with autoencoder or adversarial learning. Previous segmentation methods have also explored using autoencoder in segmentation. Oktay et al. [31] and Tong et al. [43] only trained the autoencoder with ground-truth organ shapes. The segmentation network is then trained to minimize the distance between the latent features from autoencoder of predicted shapes and ground truth shapes, where the parameters of autoencoder are fixed. However, since the autoencoder has not been trained with the predicted organ shapes by the segmentation network ever, when the estimated organ shapes are close to the ground-truth shapes, the latent codes of autoencoder alone cannot distinguish these two, thus the autoencoder cannot provide effective guidance to further regularize the segmentation network. Moreover, in Tong et al. [43], the autoencoder is trained with whole images. The autoencoder training also faces the extremely unbalanced data, thus making the training of small organs’ shape representations ineffective. In our approach, we adopt the autoencoder in the small-organ branch, thus avoid the imbalanced class problem.

Several previous works [50, 16] adopted adversarial training in segmentation networks using the objective function in the following way,

where $x$ is the input image, $y$ is its corresponding ground-truth shape, $G$ is the segmentation network, $G(x)$ is the predicted organ shape, $D$ is the discriminator network that tries to distinguish its input to be real or fake (i.e. ground truth label or network predicted mask). However, $D$ was previously designed as a classifier while we adopted the shape autoencoder for the min-max game.

The proposed FocusNetv2 was evaluated on two datasets of HaN CT images. The first dataset is a self-collected dataset, denoted as our dataset. Our dataset consists of 1,164 collected CT scans of patients with nasopharyngeal carcinoma. 22 OARs to be considered in HaN radiotherapy treatment planning are delineated in each scan, including (left and right) eyes, (left and right) lens, (left and right) optic nerves, optic chiasm, pituitary, brain stem, (left and right) temporal lobes, spinal cord, (left and right) parotid glands, (left and right) inner ears, (left and right) middle ears, (left and right) temporomandibular joints, and (left and right) mandible. Ground truth annotations of each case were provided by senior doctors with hundreds of cases of annotation experience, with each structure being segmented by the same annotator and reviewed by the other one. The left and right lens, left and right optic nerves, optic chiasm, pituitary are defined as small organs due to their small volume and complex anatomical structures. The CT scans have anisotropic voxel spacing ranging from 0.78mm to 1.25mm and inter-slice thickness ranging from 2.7mm to 3.5mm. All scans are resampled to $1\times 1\times 3$ mm for further processing. We randomly shuffle our dataset and select 1,044 samples for training and 120 samples for testing.

For comparison with state-of-the-art methods on HaN OAR segmentation, we evaluate the proposed FocusNetv2 on a public dataset, the MICCAI Head and Neck Auto Segmentation Challenge 2015 dataset (denoted as MICCAI’15 dataset). This dataset is also known as a Public Domain Database for Computational Anatomy (PDDCA) provided and maintained by Harvard Medical School. The dataset includes multiple image studies from patients with stage III or IV squamous cell carcinoma of the oropharynx, hypopharynx or larynx. It consists of 38 CT scans for training and 10 scans for testing, and has 9 organ annotations: brain stem, mandible, optic chiasm, (left and right) optic nerves, (left and right) parotid glands and (left and right) submandibular glands, where optic chiasm, left and right optic nerves are defined as small organs. The delineation of structures is based on the protocol described by Radiation Therapy Oncology Group (RTOG). We resample all scans to have a voxel size of $1\times 1\times 2.5$ mm to train our FocusNetv2, while for a fair comparison with other methods, we resampled the predicted segmentation labels by our proposed method back into the original spacing and then calculate the evaluation metrics.

Our method is implemented with PyTorch and trained on NVIDIA TITAN Xp GPUs. The segmentation networks are trained from scratch with initial weights sampled from a standard Gaussian distribution. We first train the S-Net, and then train the SOL-Net while fixing the trained parameters of S-Net. The SOS-Net is trained afterward, and is updated with the adversarial autoencoder in an alternative way. At last, we finetune the whole network for joint optimization. We use the ADAM optimizer to train the network with a learning rate 0.0005. The batch size is set as 1. For the adversarial autoencoder, it is pretrained with ground truth labels in the dataset to stabilize the adversarial training process.

The original CT image size was around $n\times 512\times 512$ with $n$ being the number of slices. Since the majority of each CT image is the background, they are centrally cropped to $n\times 240\times 240$. It would be better to use whole image volumes to train the network. However, due to GPU memory limitation, we randomly crop 40-slice chunks along the $z$-axis from CT images for each iteration when training the S-Net and the SOL-Net. One problem raised by the sliding slices strategy is that the cropping process might destroy organs’ shape for training AAE and SOS-Net. As the small organs in both datasets are mainly lens, optic nerves, optic chiasm, and pituitary. They lie in only a few adjacent slices along the $z$-axis, which can be fully included in a 40-slice cube with a large margin. Therefore, we adopt a sampling strategy when training AAE and SOS-Net. Although we sample cubes along $z$-axis with random translation, we always ensure that the cube contains all small organs with some margin. Therefore, the shapes of small organs are complete. The CT scans are cropped every 40 slices with a stride of 40 along the $z$-axis, i.e., no overlap between the crops. We then stack the segmentation result of each 40-slice crop together to obtain the final prediction. Random affine transformations (translation within 40 pixels in x and y axes, rotation within 10 degrees, and scale from 0.7 to 1.3 times) are used for data augmentation during training.

We use two evaluation metrics in this study. Dice score coefficient (DSC) measures the degree of overlap between the predicted segmentation and the ground truth segmentation with the formula, $DSC(X,Y)=\frac{2\mid X\cap Y\mid}{\mid X\mid+\mid Y\mid}$, where $X$ and $Y$ represents the voxel sets of prediction and ground-truth respectively. 95% Hausdorff Distance (95HD) is a variant of Hausdorff distance, which measures the largest distance from points in $X$ to its nearest neighbors in $Y$. The HD is calculated as the average of two directions, $HD=(d_{H}(X,Y)+d_{H}(Y,X))/2$. The 95% Hausdorff Distance can mitigate the sensitivity of outliers by calculating the 95% largest distances.

We compare our proposed method with a multi-atlas based method, where Symmetric Normalization (SyN) [2] is used as the registration method, a 3D variant of DeepLabv3+ [4] and a state-of-the-art deep learning method in HaN OAR segmentation named AnatomyNet [55].

For the multi-atlas-based method, we randomly select 9 CT scans from the training set as the atlas due to the tie and computational resouce limitation. Symmetric Normalization (SyN) [2] with its implementation in ANTs software package is used to recover the optimal affine matrix and deformable transformation field between the CT to be segmented and each atlas. 9 label maps are obtained by applying the transformation fields to atlas labels, then the final prediction is obtained by voting. DeepLabv3+ [4] is a well-known segmentation framework originally designed for 2D semantic segmentation. It uses the spatial pyramid pooling and dilated convolution and achieves state-of-the-art performance on natural image segmentation. We extend their network structure to 3D for volumetric segmentation. It was randomly initialized and trained using the same loss function as our proposed FocusNetv2. AnatomyNet [55] is designed for fast segmentation on whole CT images and has good performance compared with traditional atlas-based methods.

In this section, we provide more details about our network structure, including the segmentation network and the autoencoder.

The structure of autoencoder is shown in Table 9, which is a convolutional autoencoder. Batch normalization is also used before each ReLU layer except for the FC layer. We use trilinear upsampling layers follows convolution layers instead of transpose convolution layers in the decoder. The output of the first FC layer, which is a 512 dimentional feature vector, is the latent code of the autoencoder.

The structure of segmentation network is shown in Table 10, including S-Net, SOL-Net, and SOS-Net. For SOS-Net, the first maxpooling uses kernel size of (1,2,2) as the spacing of CT images is anisotropy, and is larger in z-axis. In the decoder, SOS-Net uses trilinear upsampling layers follow SEResBlocks to recover the spatial resolution.