Skip to content

Commit ca2ec0c

Browse files
eigenvivekeigenvivek
committed
Add 3D visualization example
1 parent e50a510 commit ca2ec0c

File tree

5 files changed

+350
-59
lines changed

5 files changed

+350
-59
lines changed

‎diffdrr/_modidx.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,9 @@
7070
'diffdrr/utils.py'),
7171
'diffdrr.utils.rotation_10d_to_quaternion': ( 'api/utils.html#rotation_10d_to_quaternion',
7272
'diffdrr/utils.py')},
73-
'diffdrr.visualization': { 'diffdrr.visualization.animate': ('api/visualization.html#animate', 'diffdrr/visualization.py'),
73+
'diffdrr.visualization': { 'diffdrr.visualization._make_camera_frustum_mesh': ( 'api/visualization.html#_make_camera_frustum_mesh',
74+
'diffdrr/visualization.py'),
75+
'diffdrr.visualization.animate': ('api/visualization.html#animate', 'diffdrr/visualization.py'),
7476
'diffdrr.visualization.drr_to_mesh': ( 'api/visualization.html#drr_to_mesh',
7577
'diffdrr/visualization.py'),
7678
'diffdrr.visualization.img_to_mesh': ( 'api/visualization.html#img_to_mesh',

‎diffdrr/visualization.py

Lines changed: 81 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -130,13 +130,14 @@ def make_fig(ground_truth):
130130
# %% ../notebooks/api/04_visualization.ipynb 10
131131
def drr_to_mesh(
132132
drr: DRR,
133+
method: str, # Either `surface_nets` or `marching_cubes`
133134
threshold: float = 300, # Min value for marching cubes (Hounsfield units)
134135
verbose: bool = True, # Display progress bars for mesh processing steps
135136
):
136137
"""
137138
Convert the CT in a DRR object into a mesh.
138139
139-
Mesh processing steps:
140+
If using marching cubes, mesh processing steps are:
140141
141142
1. Keep only largest connected components
142143
2. Smooth
@@ -150,35 +151,57 @@ def drr_to_mesh(
150151
spacing=drr.spacing,
151152
origin=(0, 0, 0),
152153
)
153-
mesh = grid.contour(
154-
isosurfaces=1,
155-
scalars=drr.volume.cpu().numpy().flatten(order="F"),
156-
rng=[threshold, torch.inf],
157-
method="marching_cubes",
158-
progress_bar=verbose,
159-
)
160154

161-
# Process the mesh
162-
mesh.extract_largest(inplace=True, progress_bar=verbose)
163-
mesh.point_data.clear()
164-
mesh.cell_data.clear()
165-
mesh.smooth_taubin(
166-
n_iter=100,
167-
feature_angle=120.0,
168-
boundary_smoothing=False,
169-
feature_smoothing=False,
170-
non_manifold_smoothing=True,
171-
normalize_coordinates=True,
172-
inplace=True,
173-
progress_bar=verbose,
174-
)
175-
mesh.decimate_pro(0.25, inplace=True, progress_bar=verbose)
176-
mesh.fill_holes(100, inplace=True, progress_bar=verbose)
177-
mesh.clean(inplace=True, progress_bar=verbose)
155+
if method == "marching_cubes":
156+
mesh = grid.contour(
157+
isosurfaces=1,
158+
scalars=drr.volume.cpu().numpy().flatten(order="F"),
159+
rng=[threshold, torch.inf],
160+
method="marching_cubes",
161+
progress_bar=verbose,
162+
)
163+
164+
# Process the mesh
165+
mesh.extract_largest(inplace=True, progress_bar=verbose)
166+
mesh.point_data.clear()
167+
mesh.cell_data.clear()
168+
mesh.smooth_taubin(
169+
n_iter=100,
170+
feature_angle=120.0,
171+
boundary_smoothing=False,
172+
feature_smoothing=False,
173+
non_manifold_smoothing=True,
174+
normalize_coordinates=True,
175+
inplace=True,
176+
progress_bar=verbose,
177+
)
178+
mesh.decimate_pro(0.25, inplace=True, progress_bar=verbose)
179+
mesh.fill_holes(100, inplace=True, progress_bar=verbose)
180+
mesh.clean(inplace=True, progress_bar=verbose)
181+
182+
elif method == "surface_nets":
183+
grid.point_data["values"] = (
184+
drr.volume.cpu().numpy().flatten(order="F") > threshold
185+
)
186+
try:
187+
mesh = grid.contour_labeled(smoothing=True, progress_bar=verbose)
188+
except AttributeError as e:
189+
raise AttributeError(
190+
f"{e}, ensure you are using pyvista>=0.43 and vtk>=9.3"
191+
)
192+
mesh.clear_cell_data()
193+
194+
else:
195+
raise ValueError(
196+
f"method must be `marching_cubes` or `surface_nets`, not {method}"
197+
)
198+
178199
return mesh
179200

180201
# %% ../notebooks/api/04_visualization.ipynb 11
181-
def img_to_mesh(drr: DRR, rotations, translations, parameterization, convention=None):
202+
def img_to_mesh(
203+
drr: DRR, rotations, translations, parameterization, convention=None, **kwargs
204+
):
182205
"""
183206
For a given pose (not batched), turn the camera and detector into a mesh.
184207
Additionally, render the DRR for the pose. Convert into a texture that
@@ -191,10 +214,14 @@ def img_to_mesh(drr: DRR, rotations, translations, parameterization, convention=
191214
img = (255.0 * img).astype(np.uint8)
192215
texture = pyvista.numpy_to_texture(img)
193216

194-
# Make a mesh for the source and detector plane
217+
# Make a mesh for the camera and the principal ray
195218
source, target = drr.detector(rotations, translations, parameterization, convention)
196-
camera = pyvista.Sphere(radius=10, center=source.squeeze().cpu().numpy())
219+
source = source.squeeze().cpu().numpy()
197220
target = target.reshape(drr.detector.height, drr.detector.width, 3).cpu().numpy()
221+
principal_ray = pyvista.Line(source, target.mean(axis=0).mean(axis=0))
222+
camera = _make_camera_frustum_mesh(source, target, size=0.125)
223+
224+
# Make a mesh for the detector plane
198225
detector = pyvista.StructuredGrid(
199226
target[..., 0],
200227
target[..., 1],
@@ -209,4 +236,29 @@ def img_to_mesh(drr: DRR, rotations, translations, parameterization, convention=
209236
inplace=True,
210237
)
211238

212-
return camera, detector, texture
239+
return camera, detector, texture, principal_ray
240+
241+
# %% ../notebooks/api/04_visualization.ipynb 12
242+
import numpy as np
243+
244+
245+
def _make_camera_frustum_mesh(source, target, size=0.125):
246+
vertices = np.stack(
247+
[
248+
source + size * (target[0, 0] - source),
249+
source + size * (target[-1, 0] - source),
250+
source + size * (target[-1, -1] - source),
251+
source + size * (target[0, -1] - source),
252+
source,
253+
]
254+
)
255+
faces = np.hstack(
256+
[
257+
[4, 0, 1, 2, 3],
258+
[3, 0, 1, 4],
259+
[3, 1, 2, 4],
260+
[3, 0, 3, 4],
261+
[3, 2, 3, 4],
262+
]
263+
)
264+
return pyvista.PolyData(vertices, faces)

‎notebooks/api/04_visualization.ipynb

Lines changed: 88 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -227,13 +227,14 @@
227227
"#| export\n",
228228
"def drr_to_mesh(\n",
229229
" drr: DRR,\n",
230+
" method: str, # Either `surface_nets` or `marching_cubes`\n",
230231
" threshold: float = 300, # Min value for marching cubes (Hounsfield units)\n",
231232
" verbose: bool = True, # Display progress bars for mesh processing steps\n",
232233
"):\n",
233234
" \"\"\"\n",
234235
" Convert the CT in a DRR object into a mesh.\n",
235236
"\n",
236-
" Mesh processing steps:\n",
237+
" If using marching cubes, mesh processing steps are:\n",
237238
"\n",
238239
" 1. Keep only largest connected components\n",
239240
" 2. Smooth\n",
@@ -247,31 +248,45 @@
247248
" spacing=drr.spacing,\n",
248249
" origin=(0, 0, 0),\n",
249250
" )\n",
250-
" mesh = grid.contour(\n",
251-
" isosurfaces=1,\n",
252-
" scalars=drr.volume.cpu().numpy().flatten(order=\"F\"),\n",
253-
" rng=[threshold, torch.inf],\n",
254-
" method=\"marching_cubes\",\n",
255-
" progress_bar=verbose,\n",
256-
" )\n",
257251
"\n",
258-
" # Process the mesh\n",
259-
" mesh.extract_largest(inplace=True, progress_bar=verbose)\n",
260-
" mesh.point_data.clear()\n",
261-
" mesh.cell_data.clear()\n",
262-
" mesh.smooth_taubin(\n",
263-
" n_iter=100,\n",
264-
" feature_angle=120.0,\n",
265-
" boundary_smoothing=False,\n",
266-
" feature_smoothing=False,\n",
267-
" non_manifold_smoothing=True,\n",
268-
" normalize_coordinates=True,\n",
269-
" inplace=True,\n",
270-
" progress_bar=verbose,\n",
271-
" )\n",
272-
" mesh.decimate_pro(0.25, inplace=True, progress_bar=verbose)\n",
273-
" mesh.fill_holes(100, inplace=True, progress_bar=verbose)\n",
274-
" mesh.clean(inplace=True, progress_bar=verbose)\n",
252+
" if method == \"marching_cubes\":\n",
253+
" mesh = grid.contour(\n",
254+
" isosurfaces=1,\n",
255+
" scalars=drr.volume.cpu().numpy().flatten(order=\"F\"),\n",
256+
" rng=[threshold, torch.inf],\n",
257+
" method=\"marching_cubes\",\n",
258+
" progress_bar=verbose,\n",
259+
" )\n",
260+
" \n",
261+
" # Process the mesh\n",
262+
" mesh.extract_largest(inplace=True, progress_bar=verbose)\n",
263+
" mesh.point_data.clear()\n",
264+
" mesh.cell_data.clear()\n",
265+
" mesh.smooth_taubin(\n",
266+
" n_iter=100,\n",
267+
" feature_angle=120.0,\n",
268+
" boundary_smoothing=False,\n",
269+
" feature_smoothing=False,\n",
270+
" non_manifold_smoothing=True,\n",
271+
" normalize_coordinates=True,\n",
272+
" inplace=True,\n",
273+
" progress_bar=verbose,\n",
274+
" )\n",
275+
" mesh.decimate_pro(0.25, inplace=True, progress_bar=verbose)\n",
276+
" mesh.fill_holes(100, inplace=True, progress_bar=verbose)\n",
277+
" mesh.clean(inplace=True, progress_bar=verbose)\n",
278+
" \n",
279+
" elif method == \"surface_nets\":\n",
280+
" grid.point_data[\"values\"] = drr.volume.cpu().numpy().flatten(order=\"F\") > threshold\n",
281+
" try:\n",
282+
" mesh = grid.contour_labeled(smoothing=True, progress_bar=verbose)\n",
283+
" except AttributeError as e:\n",
284+
" raise AttributeError(f\"{e}, ensure you are using pyvista>=0.43 and vtk>=9.3\")\n",
285+
" mesh.clear_cell_data()\n",
286+
"\n",
287+
" else:\n",
288+
" raise ValueError(f\"method must be `marching_cubes` or `surface_nets`, not {method}\")\n",
289+
" \n",
275290
" return mesh"
276291
]
277292
},
@@ -283,7 +298,7 @@
283298
"outputs": [],
284299
"source": [
285300
"#| export\n",
286-
"def img_to_mesh(drr: DRR, rotations, translations, parameterization, convention=None):\n",
301+
"def img_to_mesh(drr: DRR, rotations, translations, parameterization, convention=None, **kwargs):\n",
287302
" \"\"\"\n",
288303
" For a given pose (not batched), turn the camera and detector into a mesh.\n",
289304
" Additionally, render the DRR for the pose. Convert into a texture that\n",
@@ -296,10 +311,14 @@
296311
" img = (255.0 * img).astype(np.uint8)\n",
297312
" texture = pyvista.numpy_to_texture(img)\n",
298313
"\n",
299-
" # Make a mesh for the source and detector plane\n",
314+
" # Make a mesh for the camera and the principal ray\n",
300315
" source, target = drr.detector(rotations, translations, parameterization, convention)\n",
301-
" camera = pyvista.Sphere(radius=10, center=source.squeeze().cpu().numpy())\n",
316+
" source = source.squeeze().cpu().numpy()\n",
302317
" target = target.reshape(drr.detector.height, drr.detector.width, 3).cpu().numpy()\n",
318+
" principal_ray = pyvista.Line(source, target.mean(axis=0).mean(axis=0))\n",
319+
" camera = _make_camera_frustum_mesh(source, target, size=0.125)\n",
320+
"\n",
321+
" # Make a mesh for the detector plane\n",
303322
" detector = pyvista.StructuredGrid(\n",
304323
" target[..., 0],\n",
305324
" target[..., 1],\n",
@@ -314,7 +333,40 @@
314333
" inplace=True,\n",
315334
" )\n",
316335
"\n",
317-
" return camera, detector, texture"
336+
" return camera, detector, texture, principal_ray"
337+
]
338+
},
339+
{
340+
"cell_type": "code",
341+
"execution_count": null,
342+
"id": "fbb0de96-4efa-46ff-a337-3771c9a343e0",
343+
"metadata": {},
344+
"outputs": [],
345+
"source": [
346+
"#| exporti\n",
347+
"import numpy as np\n",
348+
"\n",
349+
"\n",
350+
"def _make_camera_frustum_mesh(source, target, size=0.125):\n",
351+
" vertices = np.stack(\n",
352+
" [\n",
353+
" source + size * (target[0, 0] - source),\n",
354+
" source + size * (target[-1, 0] - source),\n",
355+
" source + size * (target[-1, -1] - source),\n",
356+
" source + size * (target[0, -1] - source),\n",
357+
" source,\n",
358+
" ]\n",
359+
" )\n",
360+
" faces = np.hstack(\n",
361+
" [\n",
362+
" [4, 0, 1, 2, 3],\n",
363+
" [3, 0, 1, 4],\n",
364+
" [3, 1, 2, 4],\n",
365+
" [3, 0, 3, 4],\n",
366+
" [3, 2, 3, 4],\n",
367+
" ]\n",
368+
" )\n",
369+
" return pyvista.PolyData(vertices, faces)"
318370
]
319371
},
320372
{
@@ -344,6 +396,13 @@
344396
"display_name": "python3",
345397
"language": "python",
346398
"name": "python3"
399+
},
400+
"widgets": {
401+
"application/vnd.jupyter.widget-state+json": {
402+
"state": {},
403+
"version_major": 2,
404+
"version_minor": 0
405+
}
347406
}
348407
},
349408
"nbformat": 4,

‎notebooks/sidebar.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ website:
1414
- section: tutorials
1515
contents:
1616
- tutorials/introduction.ipynb
17+
- tutorials/visualization.ipynb
1718
- tutorials/metamorphasis.ipynb
1819
- tutorials/metrics.ipynb
1920
- tutorials/optimizers.ipynb

0 commit comments

Comments
 (0)