Make PoseRegressor return a RigidTrasform (#319)

eigenvivek · Aug 2, 2024 · f9ff7dc · f9ff7dc
1 parent 74bce76
commit f9ff7dc
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/diffdrr/registration.py b/diffdrr/registration.py
@@ -93,7 +93,9 @@ def forward(self, x):
         x = self.backbone(x)
         rot = self.rot_regression(x)
         xyz = self.xyz_regression(x)
-        return rot, xyz
+        return convert(
+            rot, xyz, convention=self.convention, parameterization=self.parameterization
+        )
 
 # %% ../notebooks/api/08_registration.ipynb 7
 N_ANGULAR_COMPONENTS = {

diff --git a/notebooks/api/08_registration.ipynb b/notebooks/api/08_registration.ipynb
@@ -163,7 +163,9 @@
     "        x = self.backbone(x)\n",
     "        rot = self.rot_regression(x)\n",
     "        xyz = self.xyz_regression(x)\n",
-    "        return rot, xyz"
+    "        return convert(\n",
+    "            rot, xyz, convention=self.convention, parameterization=self.parameterization\n",
+    "        )"
    ]
   },
   {