Hotfix docs (#242)

* Fix typehints

* Fix API calls

diffdrr/drr.py

@@ -104,7 +104,7 @@ def reshape_subsampled_drr(img: torch.Tensor, detector: Detector, batch_size: in
     return drr
 
 # %% ../notebooks/api/00_drr.ipynb 10
-from .pose import convert
+from .pose import RigidTransform, convert
 
 
 @patch
@@ -176,9 +176,6 @@ def set_intrinsics(
     ).to(self.volume)
 
 # %% ../notebooks/api/00_drr.ipynb 12
-from .pose import RigidTransform
-
-
 @patch
 def perspective_projection(
     self: DRR,

diffdrr/visualization.py

@@ -276,21 +276,23 @@ def labelmap_to_mesh(
 from .pose import RigidTransform
 
 
-def img_to_mesh(drr: DRR, pose: RigidTransform, **kwargs):
+def img_to_mesh(
+    drr: DRR, pose: RigidTransform, calibration: RigidTransform = None, **kwargs
+):
     """
     For a given pose (not batched), turn the camera and detector into a mesh.
     Additionally, render the DRR for the pose. Convert into a texture that
     can be applied to the detector mesh.
     """
     # Turn DRR img into a texture that can be applied to a mesh
-    img = drr(pose)
+    img = drr(pose, calibration)
     img = img.cpu().squeeze().detach().numpy()
     img = (img - img.min()) / (img.max() - img.min())
     img = (255.0 * img).astype(np.uint8)
     texture = pyvista.numpy_to_texture(img)
 
     # Make a mesh for the camera and the principal ray
-    source, target = drr.detector(pose)
+    source, target = drr.detector(pose, calibration)
     source = source.squeeze().cpu().detach().numpy()
     target = (
         target.reshape(drr.detector.height, drr.detector.width, 3)

notebooks/api/00_drr.ipynb

@@ -183,7 +183,9 @@
     "    def reshape_transform(self, img, batch_size):\n",
     "        if self.reshape:\n",
     "            if self.detector.n_subsample is None:\n",
-    "                img = img.view(batch_size, -1, self.detector.height, self.detector.width)\n",
+    "                img = img.view(\n",
+    "                    batch_size, -1, self.detector.height, self.detector.width\n",
+    "                )\n",
     "            else:\n",
     "                img = reshape_subsampled_drr(img, self.detector, batch_size)\n",
     "        return img"
@@ -221,7 +223,7 @@
    "outputs": [],
    "source": [
     "#| export\n",
-    "from diffdrr.pose import convert\n",
+    "from diffdrr.pose import RigidTransform, convert\n",
     "\n",
     "\n",
     "@patch\n",
@@ -309,9 +311,6 @@
    "outputs": [],
    "source": [
     "#| export\n",
-    "from diffdrr.pose import RigidTransform\n",
-    "\n",
-    "\n",
     "@patch\n",
     "def perspective_projection(\n",
     "    self: DRR,\n",

notebooks/api/04_visualization.ipynb

@@ -134,7 +134,7 @@
     "\n",
     "    if len(img) == 1:\n",
     "        axs = [axs]\n",
-    "    \n",
+    "\n",
     "    colors = [[int(c) for c in color[4:][:-1].split(\",\")] for color in colors]\n",
     "    masks = (img > 0).unsqueeze(-1).expand(-1, -1, -1, -1, 4)\n",
     "    masks = masks.to(torch.uint8).cpu().detach()\n",
@@ -397,21 +397,23 @@
     "from diffdrr.pose import RigidTransform\n",
     "\n",
     "\n",
-    "def img_to_mesh(drr: DRR, pose: RigidTransform, **kwargs):\n",
+    "def img_to_mesh(\n",
+    "    drr: DRR, pose: RigidTransform, calibration: RigidTransform = None, **kwargs\n",
+    "):\n",
     "    \"\"\"\n",
     "    For a given pose (not batched), turn the camera and detector into a mesh.\n",
     "    Additionally, render the DRR for the pose. Convert into a texture that\n",
     "    can be applied to the detector mesh.\n",
     "    \"\"\"\n",
     "    # Turn DRR img into a texture that can be applied to a mesh\n",
-    "    img = drr(pose)\n",
+    "    img = drr(pose, calibration)\n",
     "    img = img.cpu().squeeze().detach().numpy()\n",
     "    img = (img - img.min()) / (img.max() - img.min())\n",
     "    img = (255.0 * img).astype(np.uint8)\n",
     "    texture = pyvista.numpy_to_texture(img)\n",
     "\n",
     "    # Make a mesh for the camera and the principal ray\n",
-    "    source, target = drr.detector(pose)\n",
+    "    source, target = drr.detector(pose, calibration)\n",
     "    source = source.squeeze().cpu().detach().numpy()\n",
     "    target = (\n",
     "        target.reshape(drr.detector.height, drr.detector.width, 3)\n",