Fully convolutional neural network for 3D stroke lesion segmentation

[paper] [cite] [challenge] [code]

Abstract

The segmentation of stroke lesions is an important step for the assessment of stroke outcome. This problem can be particularly difficult if the size of the lesion is very small when compared with the size of healthy tissue. Our solution consists of what has been denoted as a U-shaped [1] or V-shaped [2] fully convolutional neural network. This type of architecture has been shown to perform well for other medical image segmentation problems. To address the fact that the lesions can be very small sized, we propose a new loss function to optimize our network. This loss function consists of the sum between the cross entropy and a “soft” dice coefficient, which is the dice coefficient calculated before binarizing the outputs. We found that this loss function leads to much faster convergence than the simple cross entropy loss or just the dice coefficient as proposed in [2]. For post processing we used a dense Conditional Random Field [3] which has been shown to improve the results of segmentation by taking advantage of the spatial information contained in the classifier’s output and the original image. Since our architecture was computationally intensive our implementation was developed using TensorFlow and Google Cloud ML engine in order to take advantage of parallelized training across multiple machines with GPUs.