Skip to content

Commit 30bf773

Browse files
eigenvivekeigenvivek
authored
Add an automatic rendering function (#354)
* Update links * Add a master 3D visualization function * Fix the labelmap vs volume flag * Mask the volume and the mask as well * Add docstring
1 parent 939b2be commit 30bf773

File tree

7 files changed

+129
-18
lines changed

7 files changed

+129
-18
lines changed

‎README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ If the 2D/3D registration capabilities are helpful, please cite our followup, [`
197197
year={2024}
198198
}
199199

200-
If you use the 3D CBCT reconstruction capabilities, please cite our followup, [`DiffVox`](https://arxiv.org/abs/2312.06358):
200+
If you use the 3D CBCT reconstruction capabilities, please cite our followup, [`DiffVox`](https://arxiv.org/abs/2411.19224):
201201

202202
@article{momeni2024voxel,
203203
title={Voxel-based Differentiable X-ray Rendering Improves Self-Supervised 3D CBCT Reconstruction},

‎diffdrr/_modidx.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@
175175
'diffdrr.utils.resample': ('api/utils.html#resample', 'diffdrr/utils.py')},
176176
'diffdrr.visualization': { 'diffdrr.visualization._make_camera_frustum_mesh': ( 'api/visualization.html#_make_camera_frustum_mesh',
177177
'diffdrr/visualization.py'),
178+
'diffdrr.visualization.add_image': ('api/visualization.html#add_image', 'diffdrr/visualization.py'),
178179
'diffdrr.visualization.animate': ('api/visualization.html#animate', 'diffdrr/visualization.py'),
179180
'diffdrr.visualization.drr_to_mesh': ( 'api/visualization.html#drr_to_mesh',
180181
'diffdrr/visualization.py'),
@@ -183,5 +184,6 @@
183184
'diffdrr.visualization.labelmap_to_mesh': ( 'api/visualization.html#labelmap_to_mesh',
184185
'diffdrr/visualization.py'),
185186
'diffdrr.visualization.plot_drr': ('api/visualization.html#plot_drr', 'diffdrr/visualization.py'),
186-
'diffdrr.visualization.plot_mask': ( 'api/visualization.html#plot_mask',
187-
'diffdrr/visualization.py')}}}
187+
'diffdrr.visualization.plot_mask': ('api/visualization.html#plot_mask', 'diffdrr/visualization.py'),
188+
'diffdrr.visualization.visualize_scene': ( 'api/visualization.html#visualize_scene',
189+
'diffdrr/visualization.py')}}}

‎diffdrr/data.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,8 @@ def read(
161161
dim=0,
162162
)
163163

164+
subject.volume.data = subject.volume.data * mask
165+
subject.mask.data = subject.mask.data * mask
164166
subject.density.data = subject.density.data * mask
165167

166168
return subject

‎diffdrr/visualization.py

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from tqdm import tqdm
1313

1414
# %% auto 0
15-
__all__ = ['plot_drr', 'plot_mask', 'animate', 'drr_to_mesh', 'labelmap_to_mesh', 'img_to_mesh']
15+
__all__ = ['plot_drr', 'plot_mask', 'animate', 'drr_to_mesh', 'labelmap_to_mesh', 'img_to_mesh', 'visualize_scene', 'add_image']
1616

1717
# %% ../notebooks/api/04_visualization.ipynb 5
1818
import torch
@@ -352,3 +352,39 @@ def _make_camera_frustum_mesh(source, target, size=0.125):
352352
]
353353
)
354354
return pyvista.PolyData(vertices, faces)
355+
356+
# %% ../notebooks/api/04_visualization.ipynb 15
357+
def visualize_scene(
358+
drr: DRR,
359+
pose: RigidTransform,
360+
labelmap: bool = False,
361+
grid: bool = True,
362+
verbose: bool = False,
363+
**kwargs
364+
):
365+
"""
366+
Given a DRR and a RigidTransform, render the 3D scene in PyVista.
367+
**kwargs are passed to drr_to_mesh.
368+
"""
369+
# Extract a mesh from the subject
370+
if labelmap:
371+
mesh = labelmap_to_mesh(drr.subject, verbose=verbose)
372+
else:
373+
mesh = drr_to_mesh(drr.subject, "surface_nets", verbose=verbose, **kwargs)
374+
375+
# Plot on a grid
376+
pl = pyvista.Plotter()
377+
pl.add_mesh(mesh)
378+
pl = add_image(drr, pose, pl)
379+
if grid:
380+
pl.show_grid()
381+
return pl
382+
383+
384+
def add_image(drr: DRR, pose: RigidTransform, pl: pyvista.Plotter):
385+
"""Add a camera to an existing scene."""
386+
camera, detector, texture, principal_ray = img_to_mesh(drr, pose)
387+
pl.add_mesh(camera, show_edges=True)
388+
pl.add_mesh(detector, texture=texture)
389+
pl.add_mesh(principal_ray, color="lime", line_width=3)
390+
return pl

‎notebooks/api/03_data.ipynb

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,8 @@
232232
" dim=0,\n",
233233
" )\n",
234234
"\n",
235+
" subject.volume.data = subject.volume.data * mask\n",
236+
" subject.mask.data = subject.mask.data * mask\n",
235237
" subject.density.data = subject.density.data * mask\n",
236238
"\n",
237239
" return subject"

‎notebooks/api/04_visualization.ipynb

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -483,6 +483,50 @@
483483
" return pyvista.PolyData(vertices, faces)"
484484
]
485485
},
486+
{
487+
"cell_type": "code",
488+
"execution_count": null,
489+
"id": "ebb5fe6b-886c-415a-8881-5387269f7950",
490+
"metadata": {},
491+
"outputs": [],
492+
"source": [
493+
"#| export\n",
494+
"def visualize_scene(\n",
495+
" drr: DRR,\n",
496+
" pose: RigidTransform,\n",
497+
" labelmap: bool = False,\n",
498+
" grid: bool = True,\n",
499+
" verbose: bool = False,\n",
500+
" **kwargs\n",
501+
"):\n",
502+
" \"\"\"\n",
503+
" Given a DRR and a RigidTransform, render the 3D scene in PyVista.\n",
504+
" **kwargs are passed to drr_to_mesh.\n",
505+
" \"\"\"\n",
506+
" # Extract a mesh from the subject\n",
507+
" if labelmap:\n",
508+
" mesh = labelmap_to_mesh(drr.subject, verbose=verbose)\n",
509+
" else:\n",
510+
" mesh = drr_to_mesh(drr.subject, \"surface_nets\", verbose=verbose, **kwargs)\n",
511+
"\n",
512+
" # Plot on a grid\n",
513+
" pl = pyvista.Plotter()\n",
514+
" pl.add_mesh(mesh)\n",
515+
" pl = add_image(drr, pose, pl)\n",
516+
" if grid:\n",
517+
" pl.show_grid()\n",
518+
" return pl\n",
519+
"\n",
520+
"\n",
521+
"def add_image(drr: DRR, pose: RigidTransform, pl: pyvista.Plotter):\n",
522+
" \"\"\"Add a camera to an existing scene.\"\"\"\n",
523+
" camera, detector, texture, principal_ray = img_to_mesh(drr, pose)\n",
524+
" pl.add_mesh(camera, show_edges=True)\n",
525+
" pl.add_mesh(detector, texture=texture)\n",
526+
" pl.add_mesh(principal_ray, color=\"lime\", line_width=3)\n",
527+
" return pl"
528+
]
529+
},
486530
{
487531
"cell_type": "code",
488532
"execution_count": null,

‎notebooks/index.ipynb

Lines changed: 39 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -48,9 +48,17 @@
4848
"source": [
4949
"## Install\n",
5050
"\n",
51+
"To install the latest stable release (**recommended**):\n",
52+
"\n",
5153
"```zsh\n",
5254
"pip install diffdrr\n",
53-
"```"
55+
"```\n",
56+
"\n",
57+
"To install the development version:\n",
58+
"\n",
59+
"```zsh\n",
60+
"git clone https://github.com/eigenvivek/DiffDRR.git --depth 1\n",
61+
"pip install -e 'DiffDRR/[dev]'"
5462
]
5563
},
5664
{
@@ -253,19 +261,36 @@
253261
"source": [
254262
"## Citing `DiffDRR`\n",
255263
"\n",
256-
"If you find `DiffDRR` useful in your work, please cite our [paper](https://doi.org/10.1007/978-3-031-23179-7_1) (or the [freely accessible arXiv version](https://arxiv.org/abs/2208.12737)):\n",
257-
"\n",
258-
"```\n",
259-
"@inproceedings{gopalakrishnanDiffDRR2022,\n",
260-
" author = {Gopalakrishnan, Vivek and Golland, Polina},\n",
261-
" title = {Fast Auto-Differentiable Digitally Reconstructed Radiographs for Solving Inverse Problems in Intraoperative Imaging},\n",
262-
" year = {2022},\n",
263-
" booktitle = {Clinical Image-based Procedures: 11th International Workshop, CLIP 2022, Held in Conjunction with MICCAI 2022, Singapore, Proceedings},\n",
264-
" series = {Lecture Notes in Computer Science},\n",
265-
" publisher = {Springer},\n",
266-
" doi = {https://doi.org/10.1007/978-3-031-23179-7_1},\n",
267-
"}\n",
268-
"```"
264+
"If you find `DiffDRR` useful in your work, please cite our\n",
265+
"[paper](https://arxiv.org/abs/2208.12737):\n",
266+
"\n",
267+
" @inproceedings{gopalakrishnan2022fast,\n",
268+
" title={Fast auto-differentiable digitally reconstructed radiographs for solving inverse problems in intraoperative imaging},\n",
269+
" author={Gopalakrishnan, Vivek and Golland, Polina},\n",
270+
" booktitle={Workshop on Clinical Image-Based Procedures},\n",
271+
" pages={1--11},\n",
272+
" year={2022},\n",
273+
" organization={Springer}\n",
274+
" }\n",
275+
"\n",
276+
"If the 2D/3D registration capabilities are helpful, please cite our followup, [`DiffPose`](https://arxiv.org/abs/2312.06358):\n",
277+
"\n",
278+
" @article{gopalakrishnan2023intraoperative,\n",
279+
" title={Intraoperative {2D/3D} image registration via differentiable {X}-ray rendering},\n",
280+
" author={Gopalakrishnan, Vivek and Dey, Neel and Golland, Polina},\n",
281+
" booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},\n",
282+
" pages={11662--11672},\n",
283+
" year={2024}\n",
284+
" }\n",
285+
"\n",
286+
"If you use the 3D CBCT reconstruction capabilities, please cite our followup, [`DiffVox`](https://arxiv.org/abs/2411.19224):\n",
287+
"\n",
288+
" @article{momeni2024voxel,\n",
289+
" title={Voxel-based Differentiable X-ray Rendering Improves Self-Supervised 3D CBCT Reconstruction},\n",
290+
" author={Momeni, Mohammadhossein and Gopalakrishnan, Vivek and Dey, Neel and Golland, Polina and Frisken, Sarah},\n",
291+
" booktitle={Machine Learning and the Physical Sciences, NeurIPS 2024},\n",
292+
" year={2024}\n",
293+
" }"
269294
]
270295
},
271296
{

0 commit comments

Comments
 (0)