From 74bce761e279f7f8d1e6cbeed6b6285d1d4766c1 Mon Sep 17 00:00:00 2001 From: Vivek Gopalakrishnan Date: Fri, 2 Aug 2024 10:20:30 -0400 Subject: [PATCH] Add intrinsic parameter resampling (#318) * Add kornia as a requirement * Add intrinsics resampling function --- diffdrr/_modidx.py | 3 +- diffdrr/utils.py | 50 +++++++++++++++++++++++++++++++- environment.yml | 1 + notebooks/api/07_utils.ipynb | 56 ++++++++++++++++++++++++++++++++++++ settings.ini | 2 +- 5 files changed, 109 insertions(+), 3 deletions(-) diff --git a/diffdrr/_modidx.py b/diffdrr/_modidx.py index d391d6f34..555f615fa 100644 --- a/diffdrr/_modidx.py +++ b/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'), diff --git a/diffdrr/utils.py b/diffdrr/utils.py index 7c3068447..0f00b8197 100644 --- a/diffdrr/utils.py +++ b/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 diff --git a/environment.yml b/environment.yml index 34884f7fc..50b2236d2 100644 --- a/environment.yml +++ b/environment.yml @@ -9,6 +9,7 @@ dependencies: - einops - torchvision - torchio + - kornia - timm - numpy - matplotlib diff --git a/notebooks/api/07_utils.ipynb b/notebooks/api/07_utils.ipynb index c63ebc77e..556d7235e 100644 --- a/notebooks/api/07_utils.ipynb +++ b/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, diff --git a/settings.ini b/settings.ini index 5fb694a1e..e060d1a7d 100644 --- a/settings.ini +++ b/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