Skip to content

Commit

Permalink
Merge pull request #169 from eigenvivek/pyvista
Browse files Browse the repository at this point in the history 
Add 3D visualization example
  • Loading branch information
@eigenvivek
eigenvivek authored Nov 25, 2023
2 parents e50a510 + 71cc74a commit 9268fd2
Show file tree
Hide file tree
Showing 5 changed files with 345 additions and 59 deletions.
4 changes: 3 additions & 1 deletion diffdrr/_modidx.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,9 @@
'diffdrr/utils.py'),
'diffdrr.utils.rotation_10d_to_quaternion': ( 'api/utils.html#rotation_10d_to_quaternion',
'diffdrr/utils.py')},
'diffdrr.visualization': { 'diffdrr.visualization.animate': ('api/visualization.html#animate', 'diffdrr/visualization.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.visualization.drr_to_mesh': ( 'api/visualization.html#drr_to_mesh',
'diffdrr/visualization.py'),
'diffdrr.visualization.img_to_mesh': ( 'api/visualization.html#img_to_mesh',
Expand Down
110 changes: 81 additions & 29 deletions diffdrr/visualization.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -130,13 +130,14 @@ def make_fig(ground_truth):
# %% ../notebooks/api/04_visualization.ipynb 10
def drr_to_mesh(
drr: DRR,
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
):
"""
Convert the CT in a DRR object into a mesh.
Mesh processing steps:
If using marching cubes, mesh processing steps are:
1. Keep only largest connected components
2. Smooth
Expand All @@ -150,35 +151,57 @@ def drr_to_mesh(
spacing=drr.spacing,
origin=(0, 0, 0),
)
mesh = grid.contour(
isosurfaces=1,
scalars=drr.volume.cpu().numpy().flatten(order="F"),
rng=[threshold, torch.inf],
method="marching_cubes",
progress_bar=verbose,
)

# Process the mesh
mesh.extract_largest(inplace=True, progress_bar=verbose)
mesh.point_data.clear()
mesh.cell_data.clear()
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.decimate_pro(0.25, inplace=True, progress_bar=verbose)
mesh.fill_holes(100, inplace=True, progress_bar=verbose)
mesh.clean(inplace=True, progress_bar=verbose)
if method == "marching_cubes":
mesh = grid.contour(
isosurfaces=1,
scalars=drr.volume.cpu().numpy().flatten(order="F"),
rng=[threshold, torch.inf],
method="marching_cubes",
progress_bar=verbose,
)

# Process the mesh
mesh.extract_largest(inplace=True, progress_bar=verbose)
mesh.point_data.clear()
mesh.cell_data.clear()
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.decimate_pro(0.25, inplace=True, progress_bar=verbose)
mesh.fill_holes(100, inplace=True, progress_bar=verbose)
mesh.clean(inplace=True, progress_bar=verbose)

elif method == "surface_nets":
grid.point_data["values"] = (
drr.volume.cpu().numpy().flatten(order="F") > threshold
)
try:
mesh = grid.contour_labeled(smoothing=True, progress_bar=verbose)
except AttributeError as e:
raise AttributeError(
f"{e}, ensure you are using pyvista>=0.43 and vtk>=9.3"
)
mesh.clear_cell_data()

else:
raise ValueError(
f"method must be `marching_cubes` or `surface_nets`, not {method}"
)

return mesh

# %% ../notebooks/api/04_visualization.ipynb 11
def img_to_mesh(drr: DRR, rotations, translations, parameterization, convention=None):
def img_to_mesh(
drr: DRR, rotations, translations, parameterization, convention=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
Expand All @@ -191,10 +214,14 @@ def img_to_mesh(drr: DRR, rotations, translations, parameterization, convention=
img = (255.0 * img).astype(np.uint8)
texture = pyvista.numpy_to_texture(img)

# Make a mesh for the source and detector plane
# Make a mesh for the camera and the principal ray
source, target = drr.detector(rotations, translations, parameterization, convention)
camera = pyvista.Sphere(radius=10, center=source.squeeze().cpu().numpy())
source = source.squeeze().cpu().numpy()
target = target.reshape(drr.detector.height, drr.detector.width, 3).cpu().numpy()
principal_ray = pyvista.Line(source, target.mean(axis=0).mean(axis=0))
camera = _make_camera_frustum_mesh(source, target, size=0.125)

# Make a mesh for the detector plane
detector = pyvista.StructuredGrid(
target[..., 0],
target[..., 1],
Expand All @@ -209,4 +236,29 @@ def img_to_mesh(drr: DRR, rotations, translations, parameterization, convention=
inplace=True,
)

return camera, detector, texture
return camera, detector, texture, principal_ray

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


def _make_camera_frustum_mesh(source, target, size=0.125):
vertices = np.stack(
[
source + size * (target[0, 0] - source),
source + size * (target[-1, 0] - source),
source + size * (target[-1, -1] - source),
source + size * (target[0, -1] - source),
source,
]
)
faces = np.hstack(
[
[4, 0, 1, 2, 3],
[3, 0, 1, 4],
[3, 1, 2, 4],
[3, 0, 3, 4],
[3, 2, 3, 4],
]
)
return pyvista.PolyData(vertices, faces)
117 changes: 88 additions & 29 deletions notebooks/api/04_visualization.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -227,13 +227,14 @@
"#| export\n",
"def drr_to_mesh(\n",
" drr: DRR,\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",
"):\n",
" \"\"\"\n",
" Convert the CT in a DRR object into a mesh.\n",
"\n",
" Mesh processing steps:\n",
" If using marching cubes, mesh processing steps are:\n",
"\n",
" 1. Keep only largest connected components\n",
" 2. Smooth\n",
Expand All @@ -247,31 +248,45 @@
" spacing=drr.spacing,\n",
" origin=(0, 0, 0),\n",
" )\n",
" mesh = grid.contour(\n",
" isosurfaces=1,\n",
" scalars=drr.volume.cpu().numpy().flatten(order=\"F\"),\n",
" rng=[threshold, torch.inf],\n",
" method=\"marching_cubes\",\n",
" progress_bar=verbose,\n",
" )\n",
"\n",
" # Process the mesh\n",
" mesh.extract_largest(inplace=True, progress_bar=verbose)\n",
" mesh.point_data.clear()\n",
" mesh.cell_data.clear()\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.decimate_pro(0.25, inplace=True, progress_bar=verbose)\n",
" mesh.fill_holes(100, inplace=True, progress_bar=verbose)\n",
" mesh.clean(inplace=True, progress_bar=verbose)\n",
" if method == \"marching_cubes\":\n",
" mesh = grid.contour(\n",
" isosurfaces=1,\n",
" scalars=drr.volume.cpu().numpy().flatten(order=\"F\"),\n",
" rng=[threshold, torch.inf],\n",
" method=\"marching_cubes\",\n",
" progress_bar=verbose,\n",
" )\n",
" \n",
" # Process the mesh\n",
" mesh.extract_largest(inplace=True, progress_bar=verbose)\n",
" mesh.point_data.clear()\n",
" mesh.cell_data.clear()\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.decimate_pro(0.25, inplace=True, progress_bar=verbose)\n",
" mesh.fill_holes(100, inplace=True, progress_bar=verbose)\n",
" mesh.clean(inplace=True, progress_bar=verbose)\n",
" \n",
" elif method == \"surface_nets\":\n",
" grid.point_data[\"values\"] = drr.volume.cpu().numpy().flatten(order=\"F\") > threshold\n",
" try:\n",
" mesh = grid.contour_labeled(smoothing=True, progress_bar=verbose)\n",
" except AttributeError as e:\n",
" raise AttributeError(f\"{e}, ensure you are using pyvista>=0.43 and vtk>=9.3\")\n",
" mesh.clear_cell_data()\n",
"\n",
" else:\n",
" raise ValueError(f\"method must be `marching_cubes` or `surface_nets`, not {method}\")\n",
" \n",
" return mesh"
]
},
Expand All @@ -283,7 +298,7 @@
"outputs": [],
"source": [
"#| export\n",
"def img_to_mesh(drr: DRR, rotations, translations, parameterization, convention=None):\n",
"def img_to_mesh(drr: DRR, rotations, translations, parameterization, convention=None, **kwargs):\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",
Expand All @@ -296,10 +311,14 @@
" img = (255.0 * img).astype(np.uint8)\n",
" texture = pyvista.numpy_to_texture(img)\n",
"\n",
" # Make a mesh for the source and detector plane\n",
" # Make a mesh for the camera and the principal ray\n",
" source, target = drr.detector(rotations, translations, parameterization, convention)\n",
" camera = pyvista.Sphere(radius=10, center=source.squeeze().cpu().numpy())\n",
" source = source.squeeze().cpu().numpy()\n",
" target = target.reshape(drr.detector.height, drr.detector.width, 3).cpu().numpy()\n",
" principal_ray = pyvista.Line(source, target.mean(axis=0).mean(axis=0))\n",
" camera = _make_camera_frustum_mesh(source, target, size=0.125)\n",
"\n",
" # Make a mesh for the detector plane\n",
" detector = pyvista.StructuredGrid(\n",
" target[..., 0],\n",
" target[..., 1],\n",
Expand All @@ -314,7 +333,40 @@
" inplace=True,\n",
" )\n",
"\n",
" return camera, detector, texture"
" return camera, detector, texture, principal_ray"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fbb0de96-4efa-46ff-a337-3771c9a343e0",
"metadata": {},
"outputs": [],
"source": [
"#| exporti\n",
"import numpy as np\n",
"\n",
"\n",
"def _make_camera_frustum_mesh(source, target, size=0.125):\n",
" vertices = np.stack(\n",
" [\n",
" source + size * (target[0, 0] - source),\n",
" source + size * (target[-1, 0] - source),\n",
" source + size * (target[-1, -1] - source),\n",
" source + size * (target[0, -1] - source),\n",
" source,\n",
" ]\n",
" )\n",
" faces = np.hstack(\n",
" [\n",
" [4, 0, 1, 2, 3],\n",
" [3, 0, 1, 4],\n",
" [3, 1, 2, 4],\n",
" [3, 0, 3, 4],\n",
" [3, 2, 3, 4],\n",
" ]\n",
" )\n",
" return pyvista.PolyData(vertices, faces)"
]
},
{
Expand Down Expand Up @@ -344,6 +396,13 @@
"display_name": "python3",
"language": "python",
"name": "python3"
},
"widgets": {
"application/vnd.jupyter.widget-state+json": {
"state": {},
"version_major": 2,
"version_minor": 0
}
}
},
"nbformat": 4,
Expand Down
1 change: 1 addition & 0 deletions notebooks/sidebar.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ website:
- section: tutorials
contents:
- tutorials/introduction.ipynb
- tutorials/visualization.ipynb
- tutorials/metamorphasis.ipynb
- tutorials/metrics.ipynb
- tutorials/optimizers.ipynb
Expand Down
Loading

0 comments on commit 9268fd2

Please sign in to comment.