GP-Unet: Lesion Detection from Weak Labels with a 3D Regression Network

TL;DR

This paper introduces GP-Unet, a 3D regression neural network that detects lesions in medical images using only global lesion count labels, enabling effective localization without detailed annotations.

Contribution

The novel GP-Unet architecture allows lesion detection from weak labels by combining regression and localization in a fully convolutional network.

Findings

Achieves 62% sensitivity in lesion detection

Outperforms thresholding and saliency-based methods by 20%

Operates effectively with only global lesion count labels

Abstract

We propose a novel convolutional neural network for lesion detection from weak labels. Only a single, global label per image - the lesion count - is needed for training. We train a regression network with a fully convolutional architecture combined with a global pooling layer to aggregate the 3D output into a scalar indicating the lesion count. When testing on unseen images, we first run the network to estimate the number of lesions. Then we remove the global pooling layer to compute localization maps of the size of the input image. We evaluate the proposed network on the detection of enlarged perivascular spaces in the basal ganglia in MRI. Our method achieves a sensitivity of 62% with on average 1.5 false positives per image. Compared with four other approaches based on intensity thresholding, saliency and class maps, our method has a 20% higher sensitivity.

Figures

Tables

Table 1. Table 1 : Results of our method in comparison to the baselines. TPR stands for true positive rate (i.e. sensitivity), FPav is the number of false positives per image in average and FDR stands for false discovery rate. The comparing methods are described in Sec. 3.0.3 .

Method	TPR	FPav	FDR
Intensities (a)	40.6	2.3	59.8
Saliency (b)	39.8	2.7	54.2
Saliency FCN (c)	18.7	3.3	70.2
Regression (d)	19.6	3.2	70.5
Regression FCN (e)	54.8	1.9	37.7
Intensities + Reg FCN (f)	62.0	1.5	31.4

Equations

$$\hat{y}=\sum _i{w}_i{h}_i+b,$$

$$\hat{y}=\sum _i{w}_i(G(f){)}_i+b=w.G(f)+b,$$

$$M=\sum _i{w}_i{f}_i,$$

Full text

1 Introduction

This paper addresses the problem of the detection of small structures in 3D images. We aim at developing a machine learning method requiring the least possible annotations for training. Several deep learning techniques [1, 2, 3] have recently been proposed for 3D segmentation. These methods use fully convolutional networks (FCN) [4] with a downsampling and upsampling path and can therefore detect small regions and fine details. Although efforts have been made to reduce the amount of annotations required with e.g, sparse annotations [1], these techniques still need pixel-wise annotations for training. Acquiring those is very time-consuming and often not feasible for large datasets. ††The final publication is available at Springer via http:://dx.doi.org/10.1007/978-3-319-66179-7 25

In [5] the authors propose 2D localization networks requiring only global image labels for their training. They aggregate the last features maps of their network into scalars thanks to a global pooling layer and can therefore compute a classification. Heatmaps can be obtained as a weighted sum of these feature maps. These heatmaps indicate which regions of the image contributed the most to the classification results. Because of the downsampling introduced by the pooling operations in the network, the heatmaps are several times smaller than the original input image, which makes it impossible to detect very small structures. In [6] a similar technique is applied to 3D CT lung data. This network splits into two branches, one with fully connected layers for classification, and the other with a global pooling layer for localization. The loss is computed as a weighted sum of these two terms. However, as in [5], this network is not suitable for the detection of small structures.

In this paper we propose a method to learn fine pixelwise detection from image-level labels. By combining global pooling with a fully convolutional architecture including upsampling steps as in the popular U-Net [7], we compute a heatmap of the size of the input. This heatmap reveals the presence of the targeted structures. During the training, unlike [5] or [6] where the authors use a classification network, the weights of the network are optimized to solve a regression task, where the objective is to predict the number of lesions.

We evaluate our method on the detection of enlarged perivascular spaces (EPVS) in the basal ganglia in MRI. EPVS is an emerging biomarker for cerebral small vessel disease. The detection of EPVS is a challenging task - even for a human observer - because of the very small size and variable shape. In [8, 9, 10] the authors propose different EPVS segmentation methods. However detection of individual EPVS has not been addressed much in the literature. Only [9] proposes an automatic method for this, but this method requires pixel-wise annotated scans for training.

2 Methods

Our method takes as input a full 3D MRI brain scan, computes a smaller smooth region of interest (ROI) and inputs it to a FCN which computes a heatmap revealing the presence of the targeted brain lesions. The FCN is trained with weak labels. Its architecture changes between training and testing but the optimized weights stay the same.

2.1 Preprocessing

Scans are first registered to MNI space. A binary mask segmenting a region of interest (ROI) in the brain is computed with standard algorithms [11]. This mask is dilated and its borders are smoothed with a Gaussian kernel. After applying the mask, we crop the image around the ROI, trim its highest values and rescale its intensities such that the input to the network is $S\in[0,1]^{h\times w\times d}$.

2.2 3D Regression Fully Convolutional Network

Fig. 1 shows the architecture of our network. It is similar to the one of 3D U-Net [1] but is less deep and has, during training, a global pooling layer [12] before the last layer. The use of this layer is detailed in the next section. Our network has $3\times 3\times 3$ convolutional layers. In the encoding path, these layers are alternated with $2\times 2\times 2$ pooling layers to downsample the image and increase the size of the receptive field. In the decoding path we use $2\times 2\times 2$ upsampling layers. We use padding. Skip connections connect the encoding and decoding paths. After each convolutional layer - except for the last one - we compute a rectified linear unit activation . The depth and number of feature maps of the network are set to fit the available GPU memory. See Fig. 1 for the architecture details.

We change the last layers of our network between training and testing. The testing step is performed with a standard fully convolutional architecture. The output is a heatmap of the size of the input. During training, we use only global image labels. To compute the loss function we need therefore to collapse the image output into a scalar. For this purpose, during training only, we introduce a global pooling layer before the last layer. In others word, we train a regression network and transform it to a segmentation network for testing (See Fig.1). We detail this below.

2.2.1 Training - Regression Network.

After the last convolutional layer, instead of stacking the voxels of each feature map and feeding them to a fully connected layer, we use a global pooling layer, which computes one pooling operation for each feature map. The resulting output of such a layer is a vector of scalars $x\in\mathbb{R}^{n}$, with $n\in\mathbb{N}$ the number of features maps of the preceding layer.

Let $f_{i}\in\mathbb{R}^{h\times w\times d}$, for $i\in\{1,2,..,n\}$, be the i-th feature map received by the global pooling layer. We can write the global pooling operation as $g(f_{i})=x_{i}$, with $x_{i}\in\mathbb{R}$. Let $G$ be the underlying mapping of the global pooling layer and $f\in\mathbb{R}^{n\times h\times w\times d}$ the vector containing the $n$ feature maps $f_{i}$. We have $G(f)\in\mathbb{R}^{n}$ and in the case of global max pooling we can write $(G(f))_{i}=\max\limits_{x,y,z}f_{i,x,y,z}$.

A convolution layer with a single output feature map $\hat{y}$ and a filter size of $1\times 1\times 1$ can be written as

[TABLE]

with $w_{i}\in\mathbb{R}$, the weights, $b$ the bias and $h_{i}$ the i-th feature map of the preceding layer. Considering the layer following the global pooling layer, we can write Eq. (1) as

[TABLE]

with . denoting the scalar product, $\hat{y}\in\mathbb{R}$, $b\in\mathbb{R}$ and $w\in\mathbb{R}^{n}$. The output of the network $\hat{y}$ is the estimated lesion count. There is no activation function applied after this last convolutional layer. To optimize the weights, we compute a mean squared loss $l=\frac{1}{m}\sum_{t=1}^{m}(\hat{y}_{t}-y_{t})^{2}$, where $\hat{y}_{t}\in\mathbb{R}$ is the predicted count, $y_{t}\in\mathbb{N}$ the ground truth for volume $S_{t}$ and $m$ the number of samples.

2.2.2 Testing - Detection Network.

During testing we remove the global pooling layer, but we do not change the weights of the network. The $h_{i}$ of Eq. (1) are the feature maps $f_{i}$ defined earlier and the bias $b^{l}$ can be omitted. We can thus rewrite Eq. (2) as

[TABLE]

with $M\in\mathbb{R}^{h\times w\times d}$ the new output of the network. $M$ is a heatmap indicating the location of the targeted lesions. Note that computation of $M$ is mathematically equivalent to the class map computation proposed in [5]. However $M$ has the same size as the input $S$ and the weights $w_{i}$ are optimized for a regression problem instead of classification.

After computing the heatmap, we need to select an appropriate threshold before we can compute a segmentation and retrieve the location of the targeted lesions. This efficiently be done by using the estimate of the lesion count $\hat{y}$ provided by the neural network: we only need to keep the global pooling layer while testing. The threshold is then selected so that the number of connected components in the thresholded heatmap is equal to the lesion count.

3 Experiments

We evaluate our method on a testing set of 30 MRI scans and compute the sensitivity, false discovery rate and average number of false positive per image.

3.0.1 Data

Our scans were acquired with a 1.5 Tesla scanner. Our dataset is a subset of the Rotterdam Scan Study [13] and contains 1642 3D PD-weighted MRI scans. Note that in our dataset PD-weighted images have a signal intensity similar to T2-weighted images, the modality commonly used to detect EPVS. The voxel resolution is $0.49\times 0.49\times 0.8\text{mm}^{3}$. Each of the scans of dataset is annotated with a single, global image label: a trained observer counted the number of EPVS in the slice of the basal ganglia showing the anterior commissure. In a random subset of $30$ scans, they marked the center of all EPVS visible in this same slice. These $30$ scans are kept for testing set. The remaining dataset is randomly split into the following subsets: $1289$ scans for training and $323$ for validation.

3.0.2 Experimental Settings

The initial ROI segmentation of the BG is computed with the subcortical segmentation of FreeSurfer [11]. For registration to MNI space we use the rigid registration implemented in Elastix [14] with the mutual information as similarity measure and default settings. The Gaussian blurring kernel has a standard deviation $\sigma=2$. The preprocessed image $S$ has a size of $168\times 128\times 84$ voxels. We trim its $1\%$ highest values. We initialize the weights of the CNN by sampling from a Gaussian distribution, use Adadelta [15] for optimization and augment the training data with randomly transformed samples. A lesion is counted as detected if the center of a connected component in the thresholded heatmap is located within a radius of $x$ pixels of the annotation. In our experiments we set $x=3$ which corresponds to the average radius of the largest lesions. We implemented our algorithms in Python with Keras and Theano libraries and ran the experiments on a Nvidia GeForce GTX 1070 GPU.

3.0.3 Baseline Methods

We compare our method with four conventional approaches, using the same preprocessing (Sec. 2.1). The first method (a) thresholds the input image based on it intensities: EPVS are among the most hyperintense structures in the basal ganlgia. Both second and third methods compute saliency maps [16] using a neural network. Saliency maps are obtained by computing the gradient of the regression score with respect to the input image. These maps can be easily computed given any classification or regression network architecture. The third method (c) uses the same network architecture as ours, the second one (b) uses a regression network without the upsampling path. Finally we compared with a method similar to [5] where we replace the original classification objective with a regression. This network is the same as ours but without the upsampling part. The thresholds are chosen based on the lesion count estimated by each network as explain in Sec. 2.

3.0.4 Results

In Table 1 we report sensitivity, false discovery rate and average false positive per image (FPav). Our method (e) detects $54.8\%$ of the lesions with on average $1.9$ false positive per image. We found that the performance could be improved further by thresholding the weighted sum of the heatmap and the original voxel intensities of the image (f). This improves the sensitivity to $62.0\%$ with $1.5$ FPav. Considering the sensitivity, our method outperforms the other baseline methods by more than $20\%$ (Table. 1). By modifying the heatmap threshold to overestimate the lesion count, we can increase the sensitivity further - $71.1\%$ - at the price of more false positive - $4.4$ in average and a false discovery rate of $60\%$ (Fig. 3).

4 Discussion

Our experiments show that using a fully convolutional network and optimizing its weights for a regression problem gives the best results on our dataset (Fig. 3 and Table 1).

The methods in [6, 5] produce coarse localization heatmaps on which it is impossible to locate small structures. In our early experiments we computed similar networks, optimized for classification. The resulting heatmaps highlight the whole ROI, which does not help to solve the detection problem. Using a regression network, the more precise output allows us to extract finer-scale information, but without an upsampling path, the results are still imprecise (see Fig. 3 method (d)). In our experiments, saliency maps [16] perform similarly to voxel intensity. By construction these maps are noisy (Fig. 3 method (b)).

Considering other works in automatic detection of EPVS, in [9] the authors compute an EPVS segmentation using a random forest classifier. They train this classifier with pixelwise ground truth annotations. They also use a 7T scanner which produces high resolution images and eases small lesion detection. They report a sensitivity of $78\%$ and a false discovery rate of $55\%$. This is better than our $71\%$ sensitivity and $60\%$ false discovery rate. However contrary to [9] which requires pixelwise annotation, our method uses only weak labels for training. In [10, 8], the authors also compute segmentations, but the evaluation is limited to a correlation with a five-category visual score.

Overall we consider that our detection results are satisfactory, considering that, during training, our method did not use any ground truth information on the location or shape of the lesions, but only their count within the whole ROI.

5 Conclusion

We presented a novel 3D neural network to detect small structures from global image-level labels. Our method combines a fully convolutional architecture with global pooling. During training the weights are optimized to solve a regression problem, while during testing the network computes lesion heatmaps of the size of the input. We demonstrated the potential of our method to detect enlarged perivascular spaces in brain MRI. We obtained a sensitivity of $62\%$ and on average $1.5$ false positives per scan. Our method outperforms four other approaches with an increase of more than $20\%$ in sensitivity. We believe this method can be applied to the detection of other brain lesions.

5.0.1 Acknowledgments

This research was funded by The Netherlands Organisation for Health Research and Development (ZonMw) Project 104003005.