Change check from rotation to orthonormal

Author: eigenvivek

diffdrr/pose.py

@@ -8,7 +8,7 @@
 import torch
 
 from einops import rearrange
-from roma import is_rotation_matrix
+from roma import is_orthonormal_matrix
 
 
 class RigidTransform(torch.nn.Module):
@@ -18,11 +18,12 @@ class RigidTransform(torch.nn.Module):
     inversion, and conversions to various representations of SE(3).
     """
 
-    def __init__(self, matrix):
+    def __init__(self, matrix, eps=1e-6):
         super().__init__()
         if matrix.dim() == 2:
             matrix = matrix.unsqueeze(0)
         self.register_buffer("matrix", matrix)
+        self.eps = eps
 
     def __len__(self):
         return len(self.matrix)
@@ -44,7 +45,7 @@ def translation(self):
         return self.matrix[..., :3, 3]
 
     def inverse(self):
-        if is_rotation_matrix(self.matrix[..., :3, :3]):
+        if is_orthonormal_matrix(self.matrix[..., :3, :3], self.eps):
             R = self.matrix[..., :3, :3]
             t = self.matrix[..., :3, 3]
             Rinv = R.mT

notebooks/api/06_pose.ipynb

@@ -92,7 +92,7 @@
     "import torch\n",
     "\n",
     "from einops import rearrange\n",
-    "from roma import is_rotation_matrix\n",
+    "from roma import is_orthonormal_matrix\n",
     "\n",
     "\n",
     "class RigidTransform(torch.nn.Module):\n",
@@ -102,11 +102,12 @@
     "    inversion, and conversions to various representations of SE(3).\n",
     "    \"\"\"\n",
     "\n",
-    "    def __init__(self, matrix):\n",
+    "    def __init__(self, matrix, eps=1e-6):\n",
     "        super().__init__()\n",
     "        if matrix.dim() == 2:\n",
     "            matrix = matrix.unsqueeze(0)\n",
     "        self.register_buffer(\"matrix\", matrix)\n",
+    "        self.eps = eps\n",
     "\n",
     "    def __len__(self):\n",
     "        return len(self.matrix)\n",
@@ -128,7 +129,7 @@
     "        return self.matrix[..., :3, 3]\n",
     "\n",
     "    def inverse(self):\n",
-    "        if is_rotation_matrix(self.matrix[..., :3, :3]):\n",
+    "        if is_orthonormal_matrix(self.matrix[..., :3, :3], self.eps):\n",
     "            R = self.matrix[..., :3, :3]\n",
     "            t = self.matrix[..., :3, 3]\n",
     "            Rinv = R.mT\n",