Deep Learning in Image Registration (DLIR) methods have been tremendously successful in image registration due to their speed and ability to incorporate weak label supervision at training time. However, DLIR methods forego many of the benefits of classical optimization-based methods. The functional nature of deep networks do not guarantee that the predicted transformation is a local minima of the registration objective, the representation of the transformation (displacement/velocity field/affine) is fixed, and the networks are not robust to domain shift. Our method aims to bridge this gap between classical and learning methods by incorporating optimization as a layer in a deep network. A deep network is trained to predict multi-scale dense feature images that are registered using a black box iterative optimization solver. This optimal warp is then used to minimize image and label alignment errors. By implicitly differentiating end-to-end through an iterative optimization solver, our learned features are registration and label-aware, and the warp functions are guaranteed to be local minima of the registration objective in the feature space. Our framework shows excellent performance on in-domain datasets, and is agnostic to domain shift such as anisotropy and varying intensity profiles. For the first time, our method allows switching between arbitrary transformation representations (free-form to diffeomorphic) at test time with zero retraining. End-to-end feature learning also facilitates interpretability of features, and out-of-the-box promptability using additional label-fidelity terms at inference.
翻译:图像配准中的深度学习方法因其速度优势及训练时整合弱标签监督的能力,在该领域取得了巨大成功。然而,此类方法舍弃了传统基于优化的配准方法的诸多优点:深度网络的功能特性无法保证预测的变换是配准目标的局部极小值;变换表示形式(位移场/速度场/仿射变换)固定不变;且网络对领域偏移缺乏鲁棒性。本研究旨在通过将优化过程作为深度网络中的一层进行整合,以弥合传统方法与学习方法之间的鸿沟。我们训练一个深度网络来预测多尺度密集特征图像,这些图像通过黑盒迭代优化求解器进行配准。随后利用所得最优形变来最小化图像与标签的对齐误差。通过迭代优化求解器进行端到端的隐式微分,我们学习的特征兼具配准与标签感知特性,且形变函数在特征空间中保证为配准目标的局部极小值。本框架在领域内数据集上表现出优异性能,并对各向异性及强度分布变化等域偏移问题具有无关性。该方法首次实现了在测试时无需重新训练即可在任意变换表示形式(从自由形变到微分同胚)间灵活切换。端到端的特征学习机制还提升了特征的可解释性,并支持在推理时通过添加标签保真度项实现即时的可提示性。