Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Render labelmaps with PyVista #214

Merged
merged 3 commits into from
Apr 11, 2024
Merged

Render labelmaps with PyVista #214

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4691c34
Add labelmap PyVista function
eigenvivek Apr 11, 2024
9e30e1d
Also canonicalize labelmaps
eigenvivek Apr 11, 2024
c125a8f
Add tutorial
eigenvivek Apr 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions diffdrr/_modidx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,4 +133,6 @@
'diffdrr/visualization.py'),
'diffdrr.visualization.img_to_mesh': ( 'api/visualization.html#img_to_mesh',
'diffdrr/visualization.py'),
'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')}}}
5 changes: 5 additions & 0 deletions diffdrr/data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,11 @@ def canonicalize(subject):
for image in subject.get_images(intensity_only=False):
image.affine = Tinv.dot(image.affine)

try:
subject.mask.affine = subject.volume.affine
except KeyError:
pass

return subject

# %% ../notebooks/api/03_data.ipynb 7
Expand Down
47 changes: 36 additions & 11 deletions diffdrr/visualization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from tqdm import tqdm

# %% auto 0
__all__ = ['plot_drr', 'animate', 'drr_to_mesh', 'img_to_mesh']
__all__ = ['plot_drr', 'animate', 'drr_to_mesh', 'labelmap_to_mesh', 'img_to_mesh']

# %% ../notebooks/api/04_visualization.ipynb 5
import torch
Expand Down Expand Up @@ -137,14 +137,13 @@ def make_fig(ground_truth):
# %% ../notebooks/api/04_visualization.ipynb 9
import pyvista
import vtk

from .drr import DRR
from torchio import Subject

vtk.vtkLogger.SetStderrVerbosity(vtk.vtkLogger.ConvertToVerbosity(-1))

# %% ../notebooks/api/04_visualization.ipynb 10
def drr_to_mesh(
drr: DRR,
subject: Subject, # torchio.Subject with a `volume` attribute
method: str, # Either `surface_nets` or `marching_cubes`
threshold: float = 300, # Min value for marching cubes (Hounsfield units)
verbose: bool = True, # Display progress bars for mesh processing steps
Expand All @@ -162,24 +161,24 @@ def drr_to_mesh(
4. Fill any holes
5. Clean (remove any redundant vertices/edges)
"""
# Turn the CT into a PyVista object and run marching cubes
# Turn the CT into a PyVista object and run surface extraction
grid = pyvista.ImageData(
dimensions=drr.volume.shape,
spacing=drr.spacing,
origin=drr.origin,
dimensions=subject.volume.shape[1:],
spacing=subject.volume.spacing,
origin=subject.volume.origin,
)

if method == "marching_cubes":
mesh = grid.contour(
isosurfaces=1,
scalars=drr.volume.cpu().numpy().flatten(order="F"),
scalars=subject.volume.data[0].cpu().numpy().flatten(order="F"),
rng=[threshold, torch.inf],
method="marching_cubes",
progress_bar=verbose,
)
elif method == "surface_nets":
grid.point_data["values"] = (
drr.volume.cpu().numpy().flatten(order="F") > threshold
subject.volume.data[0].cpu().numpy().flatten(order="F") > threshold
)
try:
mesh = grid.contour_labeled(smoothing=True, progress_bar=verbose)
Expand Down Expand Up @@ -213,6 +212,32 @@ def drr_to_mesh(
return mesh

# %% ../notebooks/api/04_visualization.ipynb 11
def labelmap_to_mesh(
subject: Subject, # torchio.Subject with a `mask` attribute
verbose: bool = True, # Display progress bars for mesh processing steps
):
# Turn the 3D labelmap into a PyVista object and run SurfaceNets
grid = pyvista.ImageData(
dimensions=subject.mask.shape[1:],
spacing=subject.mask.spacing,
origin=subject.mask.origin,
)
grid.point_data["values"] = subject.mask.data[0].numpy().flatten(order="F")
mesh = grid.contour_labeled(smoothing=True, progress_bar=verbose)
mesh.smooth_taubin(
n_iter=100,
feature_angle=120.0,
boundary_smoothing=False,
feature_smoothing=False,
non_manifold_smoothing=True,
normalize_coordinates=True,
inplace=True,
progress_bar=verbose,
)
mesh.clean(inplace=True, progress_bar=verbose)
return mesh

# %% ../notebooks/api/04_visualization.ipynb 12
from .pose import RigidTransform


Expand Down Expand Up @@ -258,7 +283,7 @@ def img_to_mesh(drr: DRR, pose: RigidTransform, **kwargs):

return camera, detector, texture, principal_ray

# %% ../notebooks/api/04_visualization.ipynb 12
# %% ../notebooks/api/04_visualization.ipynb 13
import numpy as np


Expand Down
5 changes: 5 additions & 0 deletions notebooks/api/03_data.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,11 @@
" for image in subject.get_images(intensity_only=False):\n",
" image.affine = Tinv.dot(image.affine)\n",
"\n",
" try:\n",
" subject.mask.affine = subject.volume.affine\n",
" except KeyError:\n",
" pass\n",
"\n",
" return subject"
]
},
Expand Down
51 changes: 42 additions & 9 deletions notebooks/api/04_visualization.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,8 +226,7 @@
"#| export\n",
"import pyvista\n",
"import vtk\n",
"\n",
"from diffdrr.drr import DRR\n",
"from torchio import Subject\n",
"\n",
"vtk.vtkLogger.SetStderrVerbosity(vtk.vtkLogger.ConvertToVerbosity(-1))"
]
Expand All @@ -241,7 +240,7 @@
"source": [
"#| export\n",
"def drr_to_mesh(\n",
" drr: DRR,\n",
" subject: Subject, # torchio.Subject with a `volume` attribute\n",
" method: str, # Either `surface_nets` or `marching_cubes`\n",
" threshold: float = 300, # Min value for marching cubes (Hounsfield units)\n",
" verbose: bool = True, # Display progress bars for mesh processing steps\n",
Expand All @@ -259,24 +258,24 @@
" 4. Fill any holes\n",
" 5. Clean (remove any redundant vertices/edges)\n",
" \"\"\"\n",
" # Turn the CT into a PyVista object and run marching cubes\n",
" # Turn the CT into a PyVista object and run surface extraction\n",
" grid = pyvista.ImageData(\n",
" dimensions=drr.volume.shape,\n",
" spacing=drr.spacing,\n",
" origin=drr.origin,\n",
" dimensions=subject.volume.shape[1:],\n",
" spacing=subject.volume.spacing,\n",
" origin=subject.volume.origin,\n",
" )\n",
"\n",
" if method == \"marching_cubes\":\n",
" mesh = grid.contour(\n",
" isosurfaces=1,\n",
" scalars=drr.volume.cpu().numpy().flatten(order=\"F\"),\n",
" scalars=subject.volume.data[0].cpu().numpy().flatten(order=\"F\"),\n",
" rng=[threshold, torch.inf],\n",
" method=\"marching_cubes\",\n",
" progress_bar=verbose,\n",
" )\n",
" elif method == \"surface_nets\":\n",
" grid.point_data[\"values\"] = (\n",
" drr.volume.cpu().numpy().flatten(order=\"F\") > threshold\n",
" subject.volume.data[0].cpu().numpy().flatten(order=\"F\") > threshold\n",
" )\n",
" try:\n",
" mesh = grid.contour_labeled(smoothing=True, progress_bar=verbose)\n",
Expand Down Expand Up @@ -310,6 +309,40 @@
" return mesh"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "422595b9-3f28-4c0e-b549-837c8d87430d",
"metadata": {},
"outputs": [],
"source": [
"#| export\n",
"def labelmap_to_mesh(\n",
" subject: Subject, # torchio.Subject with a `mask` attribute\n",
" verbose: bool = True, # Display progress bars for mesh processing steps\n",
"):\n",
" # Turn the 3D labelmap into a PyVista object and run SurfaceNets\n",
" grid = pyvista.ImageData(\n",
" dimensions=subject.mask.shape[1:],\n",
" spacing=subject.mask.spacing,\n",
" origin=subject.mask.origin,\n",
" )\n",
" grid.point_data[\"values\"] = subject.mask.data[0].numpy().flatten(order=\"F\")\n",
" mesh = grid.contour_labeled(smoothing=True, progress_bar=verbose)\n",
" mesh.smooth_taubin(\n",
" n_iter=100,\n",
" feature_angle=120.0,\n",
" boundary_smoothing=False,\n",
" feature_smoothing=False,\n",
" non_manifold_smoothing=True,\n",
" normalize_coordinates=True,\n",
" inplace=True,\n",
" progress_bar=verbose,\n",
" )\n",
" mesh.clean(inplace=True, progress_bar=verbose)\n",
" return mesh"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand Down
19 changes: 19 additions & 0 deletions notebooks/tutorials/mask.html
View file
Open in desktop

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion notebooks/tutorials/render.html
View file
Open in desktop

Large diffs are not rendered by default.

139 changes: 105 additions & 34 deletions notebooks/tutorials/visualization.ipynb
View file
Open in desktop

Large diffs are not rendered by default.

Loading