An Artificial Agent for Robust Image Registration
DOI:
https://doi.org/10.1609/aaai.v31i1.11230Keywords:
Image Registration, Supervised Learning, Reinforcement Learning, Deep NetworksAbstract
3-D image registration, which involves aligning two or more images, is a critical step in a variety of medical applications from diagnosis to therapy. Image registration is commonly performed by optimizing an image matching metric as a cost function. However this task is challenging due to the non-convex nature of the matching metric over the plausible registration parameter space and insufficient approches for a robust optimization. As a result, current approaches are often customized to a specific problem and sensitive to image quality and artifacts. In this paper, we propose a completely different approach to image registration, inspired by how experts perform the task. We first cast the image registration problem as a "strategic learning" process, where the goal is to find the best sequence of motion actions (e.g. up, down, etc) that yields image alignment. Within this approach, an artificial agent is learned, modeled using deep convolutional neural networks, with 3D raw image data as the input, and the next optimal action as the output. To copy with the dimensionality of the problem, we propose a greedy supervised approach for an end-to-end training, coupled with attention-driven hierarchical strategy. The resulting registration approach inherently encodes both a data-driven matching metric and an optimal registration strategy (policy). We demonstrate on two 3-D/3-D medical image registration examples with drastically different nature of challenges, that the artificial agent outperforms several state-of-the-art registration methods by a large margin in terms of both accuracy and robustness.