Add intrinsic parameter resampling (#318)

* Add kornia as a requirement

* Add intrinsics resampling function

diffdrr/_modidx.py

@@ -162,7 +162,8 @@
             'diffdrr.utils': { '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.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.parse_intrinsic_matrix': ('api/utils.html#parse_intrinsic_matrix', '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'),
                                        'diffdrr.visualization.animate': ('api/visualization.html#animate', 'diffdrr/visualization.py'),

diffdrr/utils.py

@@ -1,7 +1,7 @@
 # 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']
+__all__ = ['get_focal_length', 'get_principal_point', 'parse_intrinsic_matrix', 'make_intrinsic_matrix', 'resample']
 
 # %% ../notebooks/api/07_utils.ipynb 4
 def get_focal_length(
@@ -62,3 +62,51 @@ def make_intrinsic_matrix(
         #     [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,
+):
+    """Resample an image with new intrinsic parameters."""
+    if new_focal_len is None:
+        new_focal_len = focal_len
+    if new_delx is None:
+        new_delx = delx
+    if new_x0 is None:
+        new_x0 = x0
+    if new_y0 is None:
+        new_y0 = y0
+
+    x = img.clone()
+    _, _, height, width = x.shape
+    shape = torch.tensor([height, width])
+
+    # Translate the image
+    translation = torch.tensor([[new_x0 - x0, new_y0 - y0]]) / delx
+    x = translate(x, translation.to(x))
+
+    # Crop the image to change the focal length
+    focal_scaling = new_focal_len / focal_len
+    crop_size = (shape / focal_scaling).to(int).tolist()
+    x = center_crop(x, crop_size)
+    x = resize(x, (height, width))
+
+    # Pad the image to resize pixels
+    pixel_scaling = new_delx / delx
+    padding = (shape * (pixel_scaling - 1) / 2).to(int).tolist()
+    padding = [padding[1], padding[1], padding[0], padding[0]]
+    x = torch.nn.functional.pad(x, padding)
+    x = resize(x, (height, width))
+
+    return x

environment.yml

@@ -9,6 +9,7 @@ dependencies:
   - einops
   - torchvision
   - torchio
+  - kornia
   - timm
   - numpy
   - matplotlib

notebooks/api/07_utils.ipynb

@@ -137,6 +137,62 @@
     "    )"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ae47ec2a-0ce5-4317-8de4-47948bb4fb3f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "from kornia.geometry.transform import center_crop, resize, translate\n",
+    "\n",
+    "\n",
+    "def resample(\n",
+    "    img,\n",
+    "    focal_len,\n",
+    "    delx,\n",
+    "    x0=0,\n",
+    "    y0=0,\n",
+    "    new_focal_len=None,\n",
+    "    new_delx=None,\n",
+    "    new_x0=None,\n",
+    "    new_y0=None,\n",
+    "):\n",
+    "    \"\"\"Resample an image with new intrinsic parameters.\"\"\"\n",
+    "    if new_focal_len is None:\n",
+    "        new_focal_len = focal_len\n",
+    "    if new_delx is None:\n",
+    "        new_delx = delx\n",
+    "    if new_x0 is None:\n",
+    "        new_x0 = x0\n",
+    "    if new_y0 is None:\n",
+    "        new_y0 = y0\n",
+    "\n",
+    "    x = img.clone()\n",
+    "    _, _, height, width = x.shape\n",
+    "    shape = torch.tensor([height, width])\n",
+    "\n",
+    "    # Translate the image\n",
+    "    translation = torch.tensor([[new_x0 - x0, new_y0 - y0]]) / delx\n",
+    "    x = translate(x, translation.to(x))\n",
+    "\n",
+    "    # Crop the image to change the focal length\n",
+    "    focal_scaling = new_focal_len / focal_len\n",
+    "    crop_size = (shape / focal_scaling).to(int).tolist()\n",
+    "    x = center_crop(x, crop_size)\n",
+    "    x = resize(x, (height, width))\n",
+    "\n",
+    "    # Pad the image to resize pixels\n",
+    "    pixel_scaling = new_delx / delx\n",
+    "    padding = (shape * (pixel_scaling - 1) / 2).to(int).tolist()\n",
+    "    padding = [padding[1], padding[1], padding[0], padding[0]]\n",
+    "    x = torch.nn.functional.pad(x, padding)\n",
+    "    x = resize(x, (height, width))\n",
+    "\n",
+    "    return x"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

settings.ini

@@ -26,7 +26,7 @@ keywords = nbdev jupyter notebook python
 language = English
 status = 3
 user = eigenvivek
-requirements = matplotlib seaborn tqdm imageio fastcore 'pyvista[all]' einops torchvision scipy torchio timm numpy
+requirements = matplotlib seaborn tqdm imageio fastcore 'pyvista[all]' einops torchvision scipy torchio timm numpy kornia
 pip_requirements = torch
 conda_requirements = pytorch
 dev_requirements = nbdev black flake8 ipykernel ipywidgets jupyterlab jupyterlab_execute_time jupyterlab-code-formatter isort