Add principal point from intrinsic matrix

Author: eigenvivek

diffdrr/_modidx.py

@@ -59,6 +59,9 @@
                                'diffdrr.utils._convert_to_rotation_matrix': ( 'api/utils.html#_convert_to_rotation_matrix',
                                                                               'diffdrr/utils.py'),
                                'diffdrr.utils.convert': ('api/utils.html#convert', 'diffdrr/utils.py'),
+                               'diffdrr.utils.get_focal_length': ('api/utils.html#get_focal_length', 'diffdrr/utils.py'),
+                               'diffdrr.utils.get_principal_point': ('api/utils.html#get_principal_point', 'diffdrr/utils.py'),
+                               'diffdrr.utils.parse_intrinsic_matrix': ('api/utils.html#parse_intrinsic_matrix', 'diffdrr/utils.py'),
                                'diffdrr.utils.quaternion_adjugate_to_quaternion': ( 'api/utils.html#quaternion_adjugate_to_quaternion',
                                                                                     'diffdrr/utils.py'),
                                'diffdrr.utils.quaternion_to_quaternion_adjugate': ( 'api/utils.html#quaternion_to_quaternion_adjugate',

diffdrr/detector.py

@@ -30,6 +30,8 @@ def __init__(
         width: int,  # Width of the X-ray detector
         delx: float,  # Pixel spacing in the X-direction
         dely: float,  # Pixel spacing in the Y-direction
+        x0: float,  # Principal point X-offset
+        y0: float,  # Principal point Y-offset
         n_subsample: int | None = None,  # Number of target points to randomly sample
         reverse_x_axis: bool = False,  # If pose includes reflection (in E(3) not SE(3)), reverse x-axis
     ):
@@ -39,6 +41,8 @@ def __init__(
         self.width = width
         self.delx = delx
         self.dely = dely
+        self.x0 = x0
+        self.y0 = y0
         self.n_subsample = n_subsample
         if self.n_subsample is not None:
             self.subsamples = []
@@ -79,6 +83,10 @@ def _initialize_carm(self: Detector):
     source = source.unsqueeze(0)
     target = target.unsqueeze(0)
 
+    # Apply principal point offset
+    target[..., 2] -= self.x0
+    target[..., 1] -= self.y0
+
     if self.n_subsample is not None:
         sample = torch.randperm(self.height * self.width)[: int(self.n_subsample)]
         target = target[:, sample, :]

diffdrr/drr.py

@@ -28,6 +28,8 @@ def __init__(
         width: int
         | None = None,  # Width of the rendered DRR (if not provided, set to `height`)
         dely: float | None = None,  # Y-axis pixel size (if not provided, set to `delx`)
+        x0: float = 0.0,  # Principal point X-offset
+        y0: float = 0.0,  # Principal point Y-offset
         p_subsample: float | None = None,  # Proportion of pixels to randomly subsample
         reshape: bool = True,  # Return DRR with shape (b, 1, h, w)
         reverse_x_axis: bool = False,  # If pose includes reflection (in E(3) not SE(3)), reverse x-axis
@@ -49,6 +51,8 @@ def __init__(
             width,
             delx,
             dely,
+            x0,
+            y0,
             n_subsample=n_subsample,
             reverse_x_axis=reverse_x_axis,
         )

diffdrr/utils.py

@@ -2,7 +2,8 @@
 
 # %% auto 0
 __all__ = ['rotation_10d_to_quaternion', 'quaternion_to_rotation_10d', 'quaternion_adjugate_to_quaternion',
-           'quaternion_to_quaternion_adjugate', 'convert']
+           'quaternion_to_quaternion_adjugate', 'convert', 'get_focal_length', 'get_principal_point',
+           'parse_intrinsic_matrix']
 
 # %% ../notebooks/api/06_utils.ipynb 4
 import torch
@@ -13,9 +14,9 @@
     "euler_angles",
     "matrix",
     "quaternion",
+    "quaternion_adjugate",
     "rotation_6d",
     "rotation_10d",
-    "quaternion_adjugate",
 ]
 
 # %% ../notebooks/api/06_utils.ipynb 6
@@ -153,3 +154,37 @@ def _convert_from_rotation_matrix(matrix, parameterization, convention=None):
             f"parameterization must be in {PARAMETERIZATIONS}, not {parameterization}"
         )
     return rotation
+
+# %% ../notebooks/api/06_utils.ipynb 13
+def get_focal_length(
+    intrinsic,  # Intrinsic matrix (3 x 3 tensor)
+    delx: float,  # X-direction spacing (in units length)
+    dely: float,  # Y-direction spacing (in units length)
+) -> float:  # Focal length (in units length)
+    fx = intrinsic[0, 0]
+    fy = intrinsic[1, 1]
+    return abs((fx * delx) + (fy * delx)).item() / 2.0
+
+# %% ../notebooks/api/06_utils.ipynb 14
+def get_principal_point(
+    intrinsic,  # Intrinsic matrix (3 x 3 tensor)
+    height: int,  # Y-direction length
+    width: int,  # X-direction length
+    delx: float,  # X-direction spacing (in units length)
+    dely: float,  # Y-direction spacing (in units length)
+):
+    x0 = delx * (width / 2 - intrinsic[0, 2])
+    y0 = dely * (height / 2 - intrinsic[1, 2])
+    return x0.item(), y0.item()
+
+# %% ../notebooks/api/06_utils.ipynb 15
+def parse_intrinsic_matrix(
+    intrinsic,  # Intrinsic matrix (3 x 3 tensor)
+    height: int,  # Y-direction length
+    width: int,  # X-direction length
+    delx: float,  # X-direction spacing (in units length)
+    dely: float,  # Y-direction spacing (in units length)
+):
+    focal_length = get_focal_length(intrinsic, delx, dely)
+    x0, y0 = get_principal_point(intrinsic, height, width, delx, dely)
+    return focal_length, x0, y0

notebooks/api/00_drr.ipynb

@@ -121,6 +121,8 @@
     "        delx: float,  # X-axis pixel size\n",
     "        width: int | None = None,  # Width of the rendered DRR (if not provided, set to `height`)\n",
     "        dely: float | None = None,  # Y-axis pixel size (if not provided, set to `delx`)\n",
+    "        x0: float = 0.0,  # Principal point X-offset\n",
+    "        y0: float = 0.0,  # Principal point Y-offset\n",
     "        p_subsample: float | None = None,  # Proportion of pixels to randomly subsample\n",
     "        reshape: bool = True,  # Return DRR with shape (b, 1, h, w)\n",
     "        reverse_x_axis: bool = False,  # If pose includes reflection (in E(3) not SE(3)), reverse x-axis\n",
@@ -141,6 +143,8 @@
     "            width,\n",
     "            delx,\n",
     "            dely,\n",
+    "            x0,\n",
+    "            y0,\n",
     "            n_subsample=n_subsample,\n",
     "            reverse_x_axis=reverse_x_axis,\n",
     "        )\n",

notebooks/api/02_detector.ipynb

@@ -93,6 +93,8 @@
     "        width: int,  # Width of the X-ray detector\n",
     "        delx: float,  # Pixel spacing in the X-direction\n",
     "        dely: float,  # Pixel spacing in the Y-direction\n",
+    "        x0: float,  # Principal point X-offset\n",
+    "        y0: float,  # Principal point Y-offset\n",
     "        n_subsample: int | None = None,  # Number of target points to randomly sample\n",
     "        reverse_x_axis: bool = False,  # If pose includes reflection (in E(3) not SE(3)), reverse x-axis\n",
     "    ):\n",
@@ -102,6 +104,8 @@
     "        self.width = width\n",
     "        self.delx = delx\n",
     "        self.dely = dely\n",
+    "        self.x0 = x0\n",
+    "        self.y0 = y0\n",
     "        self.n_subsample = n_subsample\n",
     "        if self.n_subsample is not None:\n",
     "            self.subsamples = []\n",
@@ -150,6 +154,10 @@
     "    source = source.unsqueeze(0)\n",
     "    target = target.unsqueeze(0)\n",
     "\n",
+    "    # Apply principal point offset\n",
+    "    target[..., 2] -= self.x0\n",
+    "    target[..., 1] -= self.y0\n",
+    "\n",
     "    if self.n_subsample is not None:\n",
     "        sample = torch.randperm(self.height * self.width)[: int(self.n_subsample)]\n",
     "        target = target[:, sample, :]\n",

notebooks/api/06_utils.ipynb

@@ -67,9 +67,9 @@
     "    \"euler_angles\",\n",
     "    \"matrix\",\n",
     "    \"quaternion\",\n",
+    "    \"quaternion_adjugate\",\n",
     "    \"rotation_6d\",\n",
     "    \"rotation_10d\",\n",
-    "    \"quaternion_adjugate\",\n",
     "]"
    ]
   },
@@ -256,6 +256,77 @@
     "    return rotation"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "735bc47e-8fbe-4e56-8b28-3a0aa36ac6c9",
+   "metadata": {},
+   "source": [
+    "## Intrinsic matrix parsing\n",
+    "\n",
+    "From a calibrated camera's intrinsic matrix, calculate the following properties:\n",
+    "\n",
+    "- Focal length (in units length)\n",
+    "- Principal point (in units length)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "21f2cbe7-2edb-497e-ae26-2810ae383b06",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "def get_focal_length(\n",
+    "    intrinsic,  # Intrinsic matrix (3 x 3 tensor)\n",
+    "    delx: float,  # X-direction spacing (in units length)\n",
+    "    dely: float,  # Y-direction spacing (in units length)\n",
+    ") -> float:  # Focal length (in units length)\n",
+    "    fx = intrinsic[0, 0]\n",
+    "    fy = intrinsic[1, 1]\n",
+    "    return abs((fx * delx) + (fy * delx)).item() / 2.0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7f8cc2a3-7589-4c24-a75a-ec1a0f921750",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "def get_principal_point(\n",
+    "    intrinsic,  # Intrinsic matrix (3 x 3 tensor)\n",
+    "    height: int,  # Y-direction length\n",
+    "    width: int,  # X-direction length\n",
+    "    delx: float,  # X-direction spacing (in units length)\n",
+    "    dely: float,  # Y-direction spacing (in units length)\n",
+    "):\n",
+    "    x0 = delx * (width / 2 - intrinsic[0, 2])\n",
+    "    y0 = dely * (height / 2 - intrinsic[1, 2])\n",
+    "    return x0.item(), y0.item()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b26ca510-b8c3-4383-ad67-4ff9d852484d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "def parse_intrinsic_matrix(\n",
+    "    intrinsic,  # Intrinsic matrix (3 x 3 tensor)\n",
+    "    height: int,  # Y-direction length\n",
+    "    width: int,  # X-direction length\n",
+    "    delx: float,  # X-direction spacing (in units length)\n",
+    "    dely: float,  # Y-direction spacing (in units length)\n",
+    "):\n",
+    "    focal_length = get_focal_length(intrinsic, delx, dely)\n",
+    "    x0, y0 = get_principal_point(intrinsic, height, width, delx, dely)\n",
+    "    return focal_length, x0, y0"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,