Uncertainty-refined image segmentation under domain shift
A method for digital image segmentation is provided. The method comprises training a neural network for image segmentation with a labeled training dataset from a first domain, wherein a subset of nodes in the neural net are dropped out during training. The neural network receives image data from a second, different domain. A vector of N values that sum to 1 is calculated for each image element, wherein each value represents an image segmentation class. A label is assigned to each image element according to the class with the highest value in the vector. Multiple inferences are performed with active dropout layers for each image element, and an uncertainty value is generated for each image element. The label of any image element with an uncertainty value above a predefined threshold is replaced with a new label corresponding to the class with the next highest value.
1. A computer-implemented method for digital image segmentation, the method comprising:
using a number of processors to perform the steps of:
training a neural network for image segmentation with a labeled training dataset from a first domain, wherein a subset of nodes in the neural net are dropped out during training;
receiving, by the neural network, image data from a second, different domain, wherein the image data comprises a number of image elements;
calculating, by the neural network, a vector of N values that sum to 1 for each image element, wherein each of the N values represents an image segmentation class;
assigning, by the neural network, a segmentation label to each image element, wherein the segmentation label corresponds to a segmentation class with a highest value in the vector calculated for the image element;
performing, by the neural network with active dropout layers, multiple inferences for each image element;
generating, by the neural network, an uncertainty value for each image element according to the inferences; and
replacing the segmentation label of any image element with an uncertainty value above a predefined threshold with a new segmentation label corresponding to a segmentation class with a next highest value in the vector for that image element.
2. The method of claim 1 , wherein the neural network is a three-dimensional V-net convolutional neural network.
3. The method of claim 1 , wherein the image elements comprise pixels.
4. The method of claim 1 , wherein the image elements comprise voxels.
5. The method of claim 1 , wherein generating the uncertainty value for each image element further comprises taking a standard deviation over inference values for each image element.
6. The method of claim 1 , wherein different domains result from differences in at least one of:
image scanning equipment;
image element histogram value;
material composition of the imaging subject; or
image resolution.
7. The method of claim 1 , wherein replacing the segmentation labels of image elements with uncertainty values above the threshold produces a largest separation between average intensity of image elements in different segmentation classes.
8. A system for digital image segmentation, the system comprising:
a storage device configured to store program instructions; and
one or more processors operably connected to the storage device and configured to execute the program instructions to cause the system to:
train a neural network for image segmentation with a labeled training dataset from a first domain, wherein a subset of nodes in the neural net are dropped out during training;
receive, by the neural network, image data from a second, different domain, wherein the image data comprises a number of image elements;
calculate, by the neural network, a vector of N values that sum to 1 for each image element, wherein each of the N values represents an image segmentation class;
assign, by the neural network, a segmentation label to each image element, wherein the segmentation label corresponds to a segmentation class with a highest value in the vector calculated for the image element;
perform, by the neural network with active dropout layers, multiple inferences for each image element;
generate, by the neural network, an uncertainty value for each image element according to the inferences; and
replace the segmentation label of any image element with an uncertainty value above a predefined threshold with a new segmentation label corresponding to a segmentation class with a next highest value in the vector for that image element.
9. The system of claim 8 , wherein the neural network is a three-dimensional V-net convolutional neural network.
10. The system of claim 8 , wherein the image elements comprise pixels.
11. The system of claim 8 , wherein the image elements comprise voxels.
12. The system of claim 8 , wherein generating the uncertainty value for each image element further comprises taking a standard deviation over inference values for each image element.
13. The system of claim 8 , wherein different domains result from differences in at least one of:
image scanning equipment;
image element histogram value;
material composition of the imaging subject; or
image resolution.
14. The system of claim 8 , wherein replacing the segmentation labels of image elements with uncertainty values above the threshold produces a largest separation between average intensity of image elements in different segmentation classes.
15. A computer program product for digital image segmentation, the computer program product comprising:
a non-transitory computer-readable storage medium having program instructions embodied thereon to perform the steps of:
training a neural network for image segmentation with a labeled training dataset from a first domain, wherein a subset of nodes in the neural net are dropped out during training;
receiving, by the neural network, image data from a second, different domain, wherein the image data comprises a number of image elements;
calculating, by the neural network, a vector of N values that sum to 1 for each image element, wherein each of the N values represents an image segmentation class;
assigning, by the neural network, a segmentation label to each image element, wherein the segmentation label corresponds to a segmentation class with a highest value in the vector calculated for the image element;
performing, by the neural network with active dropout layers, multiple inferences for each image element;
generating, by the neural network, an uncertainty value for each image element according to the inferences; and
replacing the segmentation label of any image element with an uncertainty value above a predefined threshold with a new segmentation label corresponding to a segmentation class with a next highest value in the vector for that image element.
16. The computer program product of claim 15 , wherein the neural network is a three-dimensional V-net convolutional neural network.
17. The computer program product of claim 15 , wherein the image elements comprise pixels.
18. The computer program product of claim 15 , wherein the image elements comprise voxels.
19. The computer program product of claim 15 , wherein generating the uncertainty value for each image element further comprises taking a standard deviation over inference values for each image element.
20. The computer program product of claim 15 , wherein different domains result from differences in at least one of:
image scanning equipment;
image element histogram value;
material composition of the imaging subject; or
image resolution.