Computing CNN Loss and Gradients for Pose Estimation with Riemannian Geometry

Abstract

Pose estimation, i.e. predicting a 3D rigid transformation with respect to a fixed co-ordinate frame in, SE (3), is an omnipresent problem in medical image analysis. Deep learning methods often parameterise poses with a representation that separates rotation and translation. Available frameworks do not provide means to calculate loss on a manifold, regression is usually performed using the L2-norm independently on the rotation's and the translation's parameterisations. This is a metric for linear spaces that does not take into account the Lie group structure of SE(3). We propose a general Riemannian formulation of the pose estimation problem, and train CNNs directly on SE(3) equipped with a left-invariant Riemannian metric. The loss between the ground truth and predicted pose (elements of the manifold) is calculated as the Riemannian geodesic distance, which couples together the translation and rotation components. Network weights are updated by back-propagating the gradient with respect to the predicted pose on the tangent space of the manifold SE(3). We thoroughly evaluate the effectiveness of our loss function by comparing its performance with popular and most commonly used existing methods, on tasks such as image-based localisation and intensity-based 2D/3D registration. We also show that hyper-parameters, used in our loss function to weight the contribution between rotations and translations, can be intrinsically calculated from the dataset to achieve greater performance margins.

The Impact of Machine Learning on 2D/3D Registration for Image-Guided Interventions: A Systematic Review and Perspective

“…Such methods partially or completely skip the step of precise modeling of image formation or similarity, and instead build up the knowledge in a data-driven manner. We have identified 22 studies that describe methods for direct parameter regression ( Chou and Pizer, 2013 , 2014 ; Chou et al, 2013 ; Zhao et al, 2014 ; Mitrovi et al, 2015 ; Wu et al, 2015 ; Miao et al, 2016a ; Miao et al, 2016b ; Hou et al, 2017 ; Pei et al, 2017 ; Xie et al, 2017 ; Hou et al, 2018 ; Miao et al, 2018 ; Toth et al, 2018 ; Zhang et al, 2018 ; Zheng et al, 2018 ; Guan et al, 2019 , 2020 ; Foote et al, 2019 ; Gao et al, 2020c ; Li et al, 2020 ; Xiangqian et al, 2020 ).…”

“…As noted in Section 2.5 , most methods here limit themselves to shape estimation and assume that a global rigid alignment is either performed prior or unnecessary. The remaining 14 methods consider rigid parameter regression, and we differentiate methods that infer pose directly from the target y v ( Xiangqian et al, 2020 ; Wu et al, 2015 ; Hou et al, 2017 , 2018 ; Xie et al, 2017 ; Guan et al, 2019 , 2020 ), and methods that process both y v and ( Toth et al, 2018 ; Miao et al, 2016a , b , 2018 ; Zheng et al, 2018 ; Gao et al, 2020c ; Mitrovi et al, 2015 ), and therefore, can run iteratively. Methods that rely on the target image only are relatively straight-forward and generally train a standard feed-forward CNN architecture to regress pose on large datasets comprising of multiple independent objects or anatomies.…”

“…2 Indeed, among the seven studies that trained deep CNNs on synthetic data and evaluated in a way that allowed for comparisons between synthetic and real data performance, we found quite substantial performance drops ( Miao et al, 2018 ; Bier et al, 2019 ; Gao et al, 2020c ; Doerr et al, 2020 ; Gu et al, 2020 ; Guan et al, 2020 ; Li et al, 2020 ). Worse, three studies used different evaluation metrics in synthetic and real experiments so that comparison was not possible ( Miao et al, 2016a ; Toth et al, 2018 ; Esfandiari et al, 2021 ), and perhaps worst, ten studies that trained on synthetic data never even tested (meaningfully) on real data ( Hou et al, 2017 , 2018 ; Pei et al, 2017 ; Xie et al, 2017 ; Bier et al, 2018 ; Foote et al, 2019 ; Guan et al, 2019 ; Yang and Chen, 2019 ; Neumann et al, 2020 ; Zhang et al, 2020 ). To mitigate the negative impact of domain shift, there is growing interest in domain adaptation and generalization techniques.…”

The Impact of Machine Learning on 2D/3D Registration for Image-guided Interventions: A Systematic Review and Perspective

“…Such methods partially or completely skip the step of precise modeling of image formation or similarity, and instead build up the knowledge in a data-driven manner. We have identified 22 studies that describe methods for direct parameter regression [66,67,68,69,70,49,51,52,53,71,72,73,74,54,55,56,57,75,65,76,77,78].…”

“…As noted in Section 2.5, most methods here limit themselves to shape estimation and assume that a global rigid alignment is either performed prior or unnecessary. The remaining 14 methods consider rigid parameter regression, and we differentiate methods that infer pose directly from the target y v [66,67,69,72,71,74,73], and methods that process both y v and ŷv [68,70,75,76,78,65,77], and therefore, can run iteratively. Methods that rely on the target image only are relatively straight-forward and generally train a standard feed-forward CNN architecture to regress pose on large datasets comprising of multiple independent objects or anatomies.…”