Left atrial segmentation dataset visualization

lasc-viz.py

from pathlib import Path
import pyvista as pv
from tqdm import tqdm

# Load the data
labels_paths = list(Path('./labelsTr').glob('*.nii'))
images_paths = [Path('./imagesTr') / p.name for p in labels_paths]
images = pv.read(images_paths)
labels = pv.read(labels_paths)

# Keep the image actor for manipulation
image_actors = []

# Plot images and meshes, toggle the images' visibility with a checkbox
pl = pv.Plotter(shape=(4, 5))

for idx, (image, label) in tqdm(enumerate(zip(images, labels))):
    # Extract the surface of the ground-truth label
    surface = label.contour()
    
    # Re-center everything for synchronized viewing
    center = surface.center_of_mass()
    surface = surface.translate(-center, inplace=False)
    image = image.translate(-center, inplace=False)

    # Add the the ground-truth mesh to the scene
    pl.subplot(idx % 4, idx // 4)
    pl.add_mesh(surface, color='r', reset_camera=True)

    # Extract image slice, plot it, and store the actor for manipulation
    image_slice = image.slice(normal=[1, 0, 1], origin=[0, 0, 0])
    a = pl.add_mesh(image_slice, cmap='gray', show_scalar_bar=False, reset_camera=False)
    image_actors.append(a)

# Let all cameras move together
pl.link_views()

# Toggle image visibility
def toggle_visibility(flag):
    for actor in image_actors:
        actor.SetVisibility(flag)

pl.subplot(0, 0)
pl.add_checkbox_button_widget(toggle_visibility, value=True, position=[0, 0])

# Plot it all now
pl.show()