Enhanced convolutional neural network for image segmentation
This disclosure relates to digital image segmentation and region of interest identification. A computer implemented image segmentation method and system are particularly disclosed, including a predictive model trained based on a deep fully convolutional neural network. The model is trained using a loss function in at least one intermediate layer in addition to a loss function at the final stage of the full convolutional neural network. The predictive segmentation model trained in such a manner requires less training parameters and facilitates quicker and more accurate identification of relevant local and global features in the input image. In one implementation, the fully convolutional neural network is further supplemented with a conditional adversarial neural networks iteratively trained with the fully convolutional neural network as a discriminator measuring the quality of the predictive model generated by the fully convolutional neural network.
1. A computer image segmentation system for digital images, comprising:
a communication interface circuitry;
a database;
a predictive model repository; and
a processing circuitry in communication with the database and the predictive model repository, the processing circuitry configured to:
receive a set of training images labeled with a corresponding set of ground truth segmentation masks from the database;
establish a fully convolutional neural network comprising a multi-layer contraction convolutional neural network and an expansion convolutional neural network connected in tandem;
iteratively train the fully convolution neural network in an end-to-end manner using the set of training images and the corresponding set of ground truth segmentation masks by configuring the processing circuitry to:
forward-propagate a training image from the set of training images through the fully convolutional neural network to generate a predictive segmentation mask for the training image;
generate an end loss based on a difference between the predictive segmentation mask and a ground truth segmentation mask corresponding to the training image;
generate an intermediate loss at one of multiple intermediate layers of the multi-layer contraction convolutional neural network based on an intermediate ground truth segmentation mask downsized from a ground truth segmentation mask corresponding to the training image;
back-propagate the end loss and the intermediate loss through the fully convolutional neural network; and
minimize a sum of the end loss and the intermediate loss by adjusting a set of training parameters of the fully convolutional neural network using gradient descent;
store the iteratively trained fully convolutional neural network with the set of training parameters in the predictive model repository;
receive an unlabeled input image; and
forward-propagate the unlabeled input image through the iteratively trained convolutional neural network with the set of training parameters to generate an output segmentation mask.
2. The computer image segmentation system of claim 1 , wherein at one of the intermediate layers of the multi-layer contraction convolutional neural network, the processing circuitry is configured to:
perform convolution of a set of intermediate features sliding through an input from a preceding layer of the multi-layer contraction convolutional neural network to generate a set of feature maps;
perform a max pooling within each map to obtain a contracted set of feature maps; and
output the contracted set of feature maps to a next layer of the multi-layer contraction convolutional neural network.
3. The computer image segmentation system of claim 2 , wherein at the one of the intermediate layers of the multi-layer contraction convolutional neural network, the processing circuitry is further configured to zero-pad each of the set of feature maps to a same size as a size of the input from the preceding layer.
4. The computer image segmentation system of claim 2 , wherein the intermediate loss is generated at the one of the intermediate layers of the multi-layer contraction convolutional neural network.
5. The computer image segmentation system of claim 4 , wherein:
each of the set of ground truth segmentation masks comprises a binary mask;
the intermediate loss comprises a first portion and a second portion;
the first portion is derived from max pooling among the set of feature maps generated at the one of the intermediate layers over a portion of the set of feature maps determined by a positive portion of a ground truth segmentation mask corresponding to the training image; and
the second portion is derived from average pooling among the set of feature maps generated at the one of the intermediate layers over a portion of the set of feature maps determined by a negative portion of the ground truth segmentation mask corresponding to the training image.
6. The computer image segmentation system of claim 5 , wherein a feature map generated by max pooling the set of feature maps generated at the one of the intermediate layers and a feature map generated by average pooling the set of feature maps generated at the one of the intermediate layers are forwarded to a next layer of the one of the intermediate layers of the contraction convolutional neural network.
7. The computer image segmentation system of claim 1 , wherein the expansion convolutional neural network comprises multiple layers corresponding to the multi-layer contraction convolutional neural network.
8. The computer image segmentation system of claim 7 , wherein the expansion convolutional neural network is interconnected with the multi-layer contraction convolutional neural network at each layer by receiving feature maps from the multi-layer contraction convolutional neural network.
9. The computer image segmentation system of claim 8 , wherein the feature maps received from the multi-layer contraction convolutional neural network are concatenated with corresponding feature maps in the expansion convolutional neural network.
10. The computer image segmentation system of claim 1 , wherein each of the set of ground truth segmentation masks is a binary mask or a multi-value mask.
11. The computer image segmentation system of claim 1 , wherein the processing circuitry is further configured to iteratively train the fully convolutional neural network based on an output of a discriminator designed for classifying the set of ground truth segmentation masks and segmentation masks predicted by the full convolutional neural network.
12. The computer image segmentation system of claim 11 , wherein the discriminator comprises a multi-layer convolutional neural network.
13. The computer image segmentation system of claim 1 , wherein the processing circuitry is further configured to iteratively train the fully convolutional neural network based on an output of a discriminator designed for classifying the set of ground truth segmentation masks augmented by corresponding images and segmentation masks predicted by the full convolutional neural network and augmented by images corresponding to the predicted segmentation masks.
14. The computer image segmentation system of claim 1 , wherein the processing circuitry is further configured to:
filter the unlabeled input image using the output segmentation mask to only retain portion of image data for the unlabeled input image; and
further process the unlabeled input image based on the retained portion of the image data.
15. A method for image segmentation by a computer comprising at least one processor, the method comprising:
receiving a set of training images labeled with a corresponding set of ground truth segmentation masks;
establishing a fully convolutional neural network comprising a multi-layer contraction convolutional neural network and an expansion convolutional neural network connected in tandem;
iteratively training the full convolution neural network in an end-to-end manner using the set of training images and the corresponding set of ground truth segmentation masks by:
forward-propagating a training image from the set of training images through the fully convolutional neural network to generate a predictive segmentation mask for the training image;
generating an end loss based on a difference between the predictive segmentation mask and a ground truth segmentation mask corresponding to the training image;
generating an intermediate loss at one of multiple intermediate layers of the multi-layer contraction convolutional neural network based on an intermediate ground truth segmentation mask downsized from a ground truth segmentation mask corresponding to the training image;
back-propagating the end loss and the intermediate loss through the full convolutional neural network; and
minimizing a sum of the end loss and the intermediate loss by adjusting a set of training parameters of the fully convolutional neural network using gradient descent;
storing the iteratively trained fully convolutional neural network with the set of training parameters in a predictive model repository;
receiving an unlabeled input image; and
forward-propagating the unlabeled input image through the iteratively trained convolutional neural network with the set of training parameters to generate an output segmentation mask.
16. The method of claim 15 , further comprising, at one of the intermediate layers of the multi-layer contraction convolutional neural network:
performing convolution of a set of intermediate features sliding through an input from a preceding layer of the multi-layer contraction convolutional neural network to generate a set of feature maps;
performing a max pooling within each map to obtain a contracted set of feature maps; and
outputting the contracted set of feature maps to a next layer of the multi-layer contraction convolutional neural network.
17. The method claim 16 , further comprising, at the one of the intermediate layers of the multi-layer contraction convolutional neural network, zero padding each of the set of feature maps to a same size as a size of the input from the preceding layer.
18. The method claim 16 , wherein the intermediate loss is generated at the one of the intermediate layers of the multi-layer contraction convolutional neural network.
19. The method of claim 18 , wherein:
each of the set of ground truth segmentation masks comprises a binary mask;
the intermediate loss comprises a first portion and a second portion;
the first portion is derived from max pooling among the set of feature maps generated at the one of the intermediate layers over a portion of the set of feature maps determined by a positive portion of a ground truth segmentation mask corresponding to the training image; and
the second portion is derived from average pooling among the set of feature maps generated at the one of the intermediate layers over a portion of the set of feature maps determined by a negative portion of the ground truth segmentation mask corresponding to the training image.
20. A method for segmenting a medical image of a human organ by a computer comprising at least one processor, the method comprising:
receiving a set of training medical images labeled with a corresponding set of binary ground truth segmentation masks for the human organ;
establishing a fully convolutional neural network comprising a multi-layer contraction convolutional neural network and an expansion convolutional neural network connected in tandem;
iteratively training the full convolution neural network in an end-to-end manner using the set of training medical images and the corresponding set of binary ground truth segmentation masks by:
forward-propagating a training medical image from the set of training images through the fully convolutional neural network to generate a predictive binary segmentation mask for the training medical image;
generating an end loss based on a difference between the predictive binary segmentation mask and a binary ground truth segmentation mask corresponding to the training medical image;
generating an intermediate loss at one of multiple intermediate layers of the multi-layer contraction convolutional neural network based on an intermediate binary ground truth segmentation mask downsized from the binary ground truth segmentation mask corresponding to the training medical image;
back-propagating the end loss and the intermediate loss through the full convolutional neural network; and
minimizing a sum of the end loss and the intermediate loss by adjusting a set of training parameters of the fully convolutional neural network using gradient descent;
receiving an unlabeled medical image;
forward-propagating the unlabeled medical image through the fully convolutional neural network to generate an output binary segmentation mask;
filtering the unlabeled medical image using the output binary segmentation mask to only retain portion of image data for the unlabeled medical image; and
further processing the unlabeled medical image based on the retained portion of the image data.