Uncertainty-aware deep reinforcement learning for anatomical landmark detection in medical images

1. A method for anatomical landmark detection, the method comprising:

generating, using a first machine learning sub-network of an anatomical landmark detection agent, one or more image features for a cropped region of interest of a three-dimensional (3D) medical image, wherein the cropped ROI comprises a subset of the 3D medical image centered on a current location of the anatomical landmark detection agent within the 3D medical image;

providing to a second machine learning sub network the one or more image features for use in generating Q-values corresponding to allowable movement directions to be taken by the anatomical landmark detection agent;

generating, using at least a softmax layer of the second machine learning sub-network of the anatomical landmark detection agent for a set of six allowable movement directions associated with movement of the anatomical landmark detection agent within the 3D medical image, six corresponding discrete Q-value distributions wherein the set of six allowable movement directions comprise orthogonal image directions of the 3D medical image and wherein each discrete Q-value distribution:

comprises a predicted value distribution associated with moving the anatomical landmark detection agent from the current location within the 3D medical image, in the respective allowable movement direction;

predicting an anatomical landmark location within the 3D medical image using the plurality of discrete Q-value distributions based on moving the anatomical landmark detection agent within the 3D medical image, wherein in a given state the anatomical landmark detection agent moves in a selected allowable movement direction having a Q-value which is a maximum discrete Q-value distributions expected value over the set of six allowable movement directions; and

determining an uncertainty measure for the predicted anatomical landmark location, wherein the uncertainty measure is determined based on an average full width half maximum (FWHM) value calculated using the respective discrete Q-value distribution having the maximum expected value over the set of six allowable movement directions for the given state.

2. The method of claim 1 , wherein determining the uncertainty measure for the predicted anatomical landmark location comprises determining an average FWHM value over a plurality of inference steps performed by the anatomical landmark detection agent in an inference episode.

3. The method of claim 2 , further comprising:

determining, for each of the plurality of inference steps, a selected allowable movement direction from the set of six allowable movement directions, the selected allowable movement direction having a highest expected Q-value distribution expected value.

4. The method of claim 3 , further comprising:

determining, for each of the plurality of inference steps, the FWHM value for the Q-value distribution associated with the selected allowable movement direction; and

determining the uncertainty for the predicted anatomical landmark location based on an average FWHM of the Q-value distributions associated with the selected actions.

5. The method of claim 1 , further comprising:

identifying an out-of-distribution image prediction or identifying the predicted anatomical landmark location as a prediction error, the identifying based on the uncertainty measure determined for the predicted anatomical landmark location; and

automatically triggering an error review for the 3D medical image and the predicted anatomical landmark location.

6. The method of claim 1 , wherein the anatomical landmark detection agent is a distributional deep-Q learning network (dist-DQN) agent.

7. The method of claim 1 , wherein:

the first machine learning sub-network of the anatomical landmark detection agent is a convolutional neural network (CNN); and

the CNN generates the one or more image features for an input comprising a subset of the 3D medical image, the subset including a plurality of voxels of an image volume field of view (FOV) centered on the anatomical landmark detection agent.

8. The method of claim 1 , wherein:

the second machine learning sub-network of the anatomical landmark detection agent is a dense policy network comprising a plurality of dense layers; and

the softmax layer is provided as a final layer of the dense policy network, wherein the softmax layer encodes a discrete Q-value distribution for each allowable movement direction of the set of six allowable movement directions.

9. The method of claim 1 , wherein the 3D medical image includes a magnetic resonance (MR) image or a computed tomography (CT) image.

10. The method of claim 1 , wherein:

the 3D medical image is a three-dimensional (3D) image; and

the set of six allowable movement directions comprises six orthogonal image directions of the 3D image.

11. An apparatus for anatomical landmark detection, the apparatus comprising:

at least one memory; and

at least one processor coupled to the at least one memory, the at least one processor configured to:

generate, using a first machine learning sub-network of an anatomical landmark detection agent, one or more image features for a cropped region of interest of a three-dimensional (3D) medical image, wherein the cropped ROI comprises a subset of the 3D medical image centered on a current location of the anatomical landmark detection agent within the 3D medical image;

provide to a second machine learning sub network the one or more image features for use in generating Q-values corresponding to allowable movement directions to be taken by the anatomical landmark detection agent;

generate, using at least a softmax layer of the second machine learning sub-network of the anatomical landmark detection agent for a set of six allowable movement directions associated with movement of the anatomical landmark detection agent within the 3D medical image, six corresponding discrete Q-value distributions wherein the set of six allowable movement directions comprise orthogonal image directions of the 3D medical image and wherein each discrete Q-value distribution;

comprises a predicted value distribution associated with moving the anatomical landmark detection agent from the current location within the 3D medical image, in the respective allowable movement direction;

predict an anatomical landmark location within the 3D medical image using the plurality of discrete Q-value distributions based on moving the anatomical landmark detection agent within the 3D medical image, wherein in a given state the anatomical landmark detection agent moves in a selected allowable movement direction having a Q-value which is a maximum discrete Q-value distributions expected value over the set of six allowable movement directions; and

determine an uncertainty measure for the predicted anatomical landmark location, wherein the uncertainty measure is determined based on an average full width half maximum (FWHM) value calculated using the respective discrete Q-value distribution having the maximum expected value over the set of six allowable movement directions for the given state.

12. The apparatus of claim 11 , wherein the at least one processor is configured to determine the uncertainty measure for the predicted anatomical landmark location based on an average FWHM value determined over a plurality of inference steps associated with an inference episode of the anatomical landmark detection agent.

13. The apparatus of claim 12 , wherein the at least one processor is further configured to:

determine, for each of the plurality of inference steps, a selected allowable movement direction from the set of six allowable movement directions, the selected allowable movement direction having a highest expected Q-value distribution expected value.

14. The apparatus of claim 13 , wherein the at least one processor is further configured to:

determine, for each of the plurality of inference steps, the FWHM value for the Q-value distribution associated with the selected allowable movement direction; and

determine the uncertainty measure for the predicted anatomical landmark location based on an average FWHM value of the Q-value distributions associated with the selected movement directions.

15. The apparatus of claim 11 , wherein the at least one processor is further configured to:

identify an out-of-distribution image prediction or identify the predicted anatomical landmark location as a prediction error, the identifying based on the uncertainty measure determined for the predicted anatomical landmark location; and

automatically trigger an error review for the 3D medical image and the predicted anatomical landmark location.

16. The apparatus of claim 11 , wherein the anatomical landmark detection agent is a distributional deep-Q learning network (dist-DQN) agent.

17. The apparatus of claim 11 , wherein:

the first machine learning sub-network of the anatomical landmark detection agent is a convolutional neural network (CNN); and

the CNN generates the one or more image features for an input comprising a subset of the 3D medical image, the subset including a plurality of voxels of an image volume field of view (FOV) centered on the anatomical landmark detection agent.

18. The apparatus of claim 11 , wherein:

the second machine learning sub-network of the anatomical landmark detection agent is a dense policy network comprising a plurality of dense layers; and

the softmax layer is provided as a final layer of the dense policy network, wherein the softmax layer encodes a discrete Q-value distribution for each allowable movement direction of the set of six allowable movement directions.

19. The apparatus of claim 11 , wherein the 3D medical image includes a magnetic resonance (MR) image or a computed tomography (CT) image.

20. The apparatus of claim 11 , wherein:

The 3D medical image is a three-dimensional (3D) image; and

the set of six allowable movement directions comprises six orthogonal image directions of the 3D image.