Fix circular imports

eigenvivek · Dec 12, 2024 · b94b38e · b94b38e
1 parent 9eb5574
commit b94b38e
Show file tree

Hide file tree

Showing 5 changed files with 190 additions and 191 deletions.
diff --git a/diffdrr/_modidx.py b/diffdrr/_modidx.py
@@ -22,7 +22,13 @@
                                   'diffdrr.detector.Detector.reorient': ('api/detector.html#detector.reorient', 'diffdrr/detector.py'),
                                   'diffdrr.detector.Detector.sdd': ('api/detector.html#detector.sdd', 'diffdrr/detector.py'),
                                   'diffdrr.detector.Detector.x0': ('api/detector.html#detector.x0', 'diffdrr/detector.py'),
-                                  'diffdrr.detector.Detector.y0': ('api/detector.html#detector.y0', 'diffdrr/detector.py')},
+                                  'diffdrr.detector.Detector.y0': ('api/detector.html#detector.y0', 'diffdrr/detector.py'),
+                                  'diffdrr.detector.get_focal_length': ('api/detector.html#get_focal_length', 'diffdrr/detector.py'),
+                                  'diffdrr.detector.get_principal_point': ('api/detector.html#get_principal_point', 'diffdrr/detector.py'),
+                                  'diffdrr.detector.make_intrinsic_matrix': ( 'api/detector.html#make_intrinsic_matrix',
+                                                                              'diffdrr/detector.py'),
+                                  'diffdrr.detector.parse_intrinsic_matrix': ( 'api/detector.html#parse_intrinsic_matrix',
+                                                                               'diffdrr/detector.py')},
             'diffdrr.drr': { 'diffdrr.drr.DRR': ('api/drr.html#drr', 'diffdrr/drr.py'),
                              'diffdrr.drr.DRR.__init__': ('api/drr.html#drr.__init__', 'diffdrr/drr.py'),
                              'diffdrr.drr.DRR.affine': ('api/drr.html#drr.affine', 'diffdrr/drr.py'),
@@ -167,11 +173,7 @@
                                    'diffdrr.renderers._get_voxel': ('api/renderers.html#_get_voxel', 'diffdrr/renderers.py'),
                                    'diffdrr.renderers._get_xyzs': ('api/renderers.html#_get_xyzs', 'diffdrr/renderers.py'),
                                    'diffdrr.renderers.reduce': ('api/renderers.html#reduce', 'diffdrr/renderers.py')},
-            'diffdrr.utils': { 'diffdrr.utils.get_focal_length': ('api/utils.html#get_focal_length', 'diffdrr/utils.py'),
-                               'diffdrr.utils.get_pinhole_camera': ('api/utils.html#get_pinhole_camera', 'diffdrr/utils.py'),
-                               'diffdrr.utils.get_principal_point': ('api/utils.html#get_principal_point', 'diffdrr/utils.py'),
-                               'diffdrr.utils.make_intrinsic_matrix': ('api/utils.html#make_intrinsic_matrix', 'diffdrr/utils.py'),
-                               'diffdrr.utils.parse_intrinsic_matrix': ('api/utils.html#parse_intrinsic_matrix', 'diffdrr/utils.py'),
+            'diffdrr.utils': { 'diffdrr.utils.get_pinhole_camera': ('api/utils.html#get_pinhole_camera', 'diffdrr/utils.py'),
                                'diffdrr.utils.resample': ('api/utils.html#resample', 'diffdrr/utils.py')},
             'diffdrr.visualization': { 'diffdrr.visualization._make_camera_frustum_mesh': ( 'api/visualization.html#_make_camera_frustum_mesh',
                                                                                             'diffdrr/visualization.py'),
@@ -184,6 +186,8 @@
                                        'diffdrr.visualization.labelmap_to_mesh': ( 'api/visualization.html#labelmap_to_mesh',
                                                                                    'diffdrr/visualization.py'),
                                        'diffdrr.visualization.plot_drr': ('api/visualization.html#plot_drr', 'diffdrr/visualization.py'),
+                                       'diffdrr.visualization.plot_img_and_mask': ( 'api/visualization.html#plot_img_and_mask',
+                                                                                    'diffdrr/visualization.py'),
                                        'diffdrr.visualization.plot_mask': ('api/visualization.html#plot_mask', 'diffdrr/visualization.py'),
                                        'diffdrr.visualization.visualize_scene': ( 'api/visualization.html#visualize_scene',
                                                                                   'diffdrr/visualization.py')}}}
diff --git a/diffdrr/detector.py b/diffdrr/detector.py
@@ -8,11 +8,10 @@
 from torch.nn.functional import normalize
 
 # %% auto 0
-__all__ = ['Detector']
+__all__ = ['Detector', 'get_focal_length', 'get_principal_point', 'parse_intrinsic_matrix', 'make_intrinsic_matrix']
 
 # %% ../notebooks/api/02_detector.ipynb 5
 from .pose import RigidTransform
-from .utils import make_intrinsic_matrix
 
 
 class Detector(torch.nn.Module):
@@ -157,3 +156,55 @@ def forward(self: Detector, extrinsic: RigidTransform, calibration: RigidTransfo
     source = pose(self.source)
     target = pose(target)
     return source, target
+
+# %% ../notebooks/api/02_detector.ipynb 9
+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 * dely)).item() / 2.0
+
+# %% ../notebooks/api/02_detector.ipynb 10
+def get_principal_point(
+    intrinsic,  # Intrinsic matrix (3 x 3 tensor)
+    height: int,  # Y-direction length (in units pixels)
+    width: int,  # X-direction length (in units pixels)
+    delx: float,  # X-direction spacing (in units length)
+    dely: float,  # Y-direction spacing (in units length)
+):
+    x0 = delx * (intrinsic[0, 2] - width / 2)
+    y0 = dely * (intrinsic[1, 2] - height / 2)
+    return x0.item(), y0.item()
+
+# %% ../notebooks/api/02_detector.ipynb 11
+def parse_intrinsic_matrix(
+    intrinsic,  # Intrinsic matrix (3 x 3 tensor)
+    height: int,  # Y-direction length (in units pixels)
+    width: int,  # X-direction length (in units pixels)
+    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/02_detector.ipynb 12
+def make_intrinsic_matrix(
+    sdd: float,  # Source-to-detector distance (in units length)
+    delx: float,  # X-direction spacing (in units length / pixel)
+    dely: float,  # Y-direction spacing (in units length / pixel)
+    height: int,  # Y-direction length (in units pixels)
+    width: int,  # X-direction length (in units pixels)
+    x0: float = 0.0,  # Principal point x-coordinate (in units length)
+    y0: float = 0.0,  # Principal point y-coordinate (in units length)
+):
+    return torch.tensor(
+        [
+            [sdd / delx, 0.0, x0 / delx + width / 2],
+            [0.0, sdd / dely, y0 / dely + height / 2],
+            [0.0, 0.0, 1.0],
+        ]
+    )
diff --git a/diffdrr/utils.py b/diffdrr/utils.py
@@ -1,84 +1,24 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/api/07_utils.ipynb.
 
 # %% auto 0
-__all__ = ['get_focal_length', 'get_principal_point', 'parse_intrinsic_matrix', 'make_intrinsic_matrix', 'resample',
-           'get_pinhole_camera']
+__all__ = ['resample', 'get_pinhole_camera']
 
 # %% ../notebooks/api/07_utils.ipynb 4
-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 * dely)).item() / 2.0
-
-# %% ../notebooks/api/07_utils.ipynb 5
-def get_principal_point(
-    intrinsic,  # Intrinsic matrix (3 x 3 tensor)
-    height: int,  # Y-direction length (in units pixels)
-    width: int,  # X-direction length (in units pixels)
-    delx: float,  # X-direction spacing (in units length)
-    dely: float,  # Y-direction spacing (in units length)
-):
-    x0 = delx * (intrinsic[0, 2] - width / 2)
-    y0 = dely * (intrinsic[1, 2] - height / 2)
-    return x0.item(), y0.item()
-
-# %% ../notebooks/api/07_utils.ipynb 6
-def parse_intrinsic_matrix(
-    intrinsic,  # Intrinsic matrix (3 x 3 tensor)
-    height: int,  # Y-direction length (in units pixels)
-    width: int,  # X-direction length (in units pixels)
-    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/07_utils.ipynb 7
 import torch
-
-
-def make_intrinsic_matrix(
-    sdd: float,  # Source-to-detector distance (in units length)
-    delx: float,  # X-direction spacing (in units length / pixel)
-    dely: float,  # Y-direction spacing (in units length / pixel)
-    height: int,  # Y-direction length (in units pixels)
-    width: int,  # X-direction length (in units pixels)
-    x0: float = 0.0,  # Principal point x-coordinate (in units length)
-    y0: float = 0.0,  # Principal point y-coordinate (in units length)
-):
-    return torch.tensor(
-        [
-            [sdd / delx, 0.0, x0 / delx + width / 2],
-            [0.0, sdd / dely, y0 / dely + height / 2],
-            [0.0, 0.0, 1.0],
-        ]
-        # [
-        #     [sdd / delx, 0.0, -x0 / delx + width / 2],
-        #     [0.0, sdd / dely, -y0 / dely + height / 2],
-        #     [0.0, 0.0, 1.0],
-        # ]
-    )
-
-# %% ../notebooks/api/07_utils.ipynb 8
 from kornia.geometry.transform import center_crop, resize, translate
 
 
 def resample(
-    img,
-    focal_len,
-    delx,
-    x0=0,
-    y0=0,
-    new_focal_len=None,
-    new_delx=None,
-    new_x0=None,
-    new_y0=None,
-):
+    img: torch.Tensor,
+    focal_len: float,
+    delx: float,
+    x0: float = 0,
+    y0: float = 0,
+    new_focal_len: float = None,
+    new_delx: float = None,
+    new_x0: float = None,
+    new_y0: float = None,
+) -> torch.Tensor:
     """Resample an image with new intrinsic parameters."""
     if new_focal_len is None:
         new_focal_len = focal_len
@@ -112,14 +52,15 @@ def resample(
 
     return x
 
-# %% ../notebooks/api/07_utils.ipynb 10
+# %% ../notebooks/api/07_utils.ipynb 6
 from kornia.geometry.calibration import solve_pnp_dlt
 from kornia.geometry.camera.pinhole import PinholeCamera
 
+from .drr import DRR
 from .pose import RigidTransform
 
 
-def get_pinhole_camera(drr, pose: RigidTransform) -> PinholeCamera:
+def get_pinhole_camera(drr: DRR, pose: RigidTransform) -> PinholeCamera:
     # Move everything to CPU and use double precision
     drr = drr.to(device="cpu", dtype=torch.float64)
     pose = pose.to(device="cpu", dtype=torch.float64)

diff --git a/notebooks/api/02_detector.ipynb b/notebooks/api/02_detector.ipynb
@@ -67,7 +67,6 @@
    "source": [
     "#| export\n",
     "from diffdrr.pose import RigidTransform\n",
-    "from diffdrr.utils import make_intrinsic_matrix\n",
     "\n",
     "\n",
     "class Detector(torch.nn.Module):\n",
@@ -230,6 +229,102 @@
     "    return source, target"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "f4558157-a060-4add-b4ca-22600a26232d",
+   "metadata": {},
+   "source": [
+    "## Intrinsic matrix parsing\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": "4c0f02f6-c27e-4bdc-a204-31ba5c9f73de",
+   "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 * dely)).item() / 2.0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a3535bdf-b819-4c42-9624-00d101b29ded",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "def get_principal_point(\n",
+    "    intrinsic,  # Intrinsic matrix (3 x 3 tensor)\n",
+    "    height: int,  # Y-direction length (in units pixels)\n",
+    "    width: int,  # X-direction length (in units pixels)\n",
+    "    delx: float,  # X-direction spacing (in units length)\n",
+    "    dely: float,  # Y-direction spacing (in units length)\n",
+    "):\n",
+    "    x0 = delx * (intrinsic[0, 2] - width / 2)\n",
+    "    y0 = dely * (intrinsic[1, 2] - height / 2)\n",
+    "    return x0.item(), y0.item()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "750cb0fe-c96a-4c76-a2cd-51a74fdc6b05",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "def parse_intrinsic_matrix(\n",
+    "    intrinsic,  # Intrinsic matrix (3 x 3 tensor)\n",
+    "    height: int,  # Y-direction length (in units pixels)\n",
+    "    width: int,  # X-direction length (in units pixels)\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,
+   "id": "4e9f01cb-1dbc-4818-8521-e6785c101a82",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "def make_intrinsic_matrix(\n",
+    "    sdd: float,  # Source-to-detector distance (in units length)\n",
+    "    delx: float,  # X-direction spacing (in units length / pixel)\n",
+    "    dely: float,  # Y-direction spacing (in units length / pixel)\n",
+    "    height: int,  # Y-direction length (in units pixels)\n",
+    "    width: int,  # X-direction length (in units pixels)\n",
+    "    x0: float = 0.0,  # Principal point x-coordinate (in units length)\n",
+    "    y0: float = 0.0,  # Principal point y-coordinate (in units length)\n",
+    "):\n",
+    "    return torch.tensor(\n",
+    "        [\n",
+    "            [sdd / delx, 0.0, x0 / delx + width / 2],\n",
+    "            [0.0, sdd / dely, y0 / dely + height / 2],\n",
+    "            [0.0, 0.0, 1.0],\n",
+    "        ]\n",
+    "    )"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,