Improving Probabilistic Image Registration via Reinforcement Learning and Uncertainty Evaluation

Abstract

One framework for probabilistic image registration involves assigning probability distributions over spatial transformations (e.g. distributions over displacement vectors at each voxel). In this paper, we propose an uncertainty measure for these distributions that examines the actual spatial displacements, thus departing from the classical Shannon entropy-based measures, which examine only the probabilities of these distributions. We show that by incorporating the proposed uncertainty measure, along with features extracted from the input images and intermediate displacement fields, we are able to more accurately predict the pointwise registration errors of an intermediate solution as estimated for a previously unseen input image pair. We utilize the predicted errors to identify regions in the image that are trustworthy and through which we refine the tentative registration solution. Results show that our proposed framework, which incorporates uncertainty estimation and registration error prediction, can improve accuracy of 3D image registrations by about 25%.