SPINNING_BUNDLE_BRAIN.py

import numpy as np
import nibabel as nib
import pyvista as pv
import imageio.v2 as imageio
import argparse
from glob import glob
import os

T1_mask_path = "/valiant02/masi/rudravg/atlases/Pandora-WhiteMatterAtlas/T1/binary_mask.nii.gz"
bundles_path = "/valiant02/masi/rudravg/atlases/Pandora-WhiteMatterAtlas/TractSeg/all_files"


def nifti_to_surface(nifti_path):
    """Load a NIfTI binary mask and return a PyVista surface mesh in world coords."""
    nii = nib.load(nifti_path)
    data = nii.get_fdata().astype(np.float32)
    affine = nii.affine

    # Build a PyVista structured grid in voxel space
    grid = pv.ImageData(dimensions=data.shape)
    grid.point_data["values"] = data.ravel(order="F")

    # Extract isosurface at 0.5 (binary mask boundary)
    surface = grid.contour([0.5], scalars="values")

    if surface.n_points == 0:
        return None

    # Transform from voxel indices to world (RAS) coordinates
    surface.transform(affine)
    surface.compute_normals(inplace=True)
    return surface


parser = argparse.ArgumentParser(description="Spinning glass brain with crossfading bundle pairs")
parser.add_argument("--bundles_path", default=bundles_path, help="Dir with vol_*.nii.gz")
parser.add_argument("--brain_mask", default=T1_mask_path, help="Binary brain mask for glass brain")
parser.add_argument("--output", default="spinning_bundles.mp4", help="Output MP4")
args = parser.parse_args()

# Config
RESOLUTION = (1920, 1088)
FPS = 60
BACKGROUND_COLOR = "white"
HOLD_TIME = 1.5        # seconds each pair is visible (hard cut, no fade)
ROTATION_SPEED = 36.0  # degrees per second

BLUE   = [0.3, 0.6, 1.0]
ORANGE = [1.0, 0.55, 0.1]

bundle_files = sorted(glob(os.path.join(args.bundles_path, "vol_*.nii.gz")))
n_bundles = len(bundle_files)
n_pairs = n_bundles // 2
print(f"Found {n_bundles} bundle masks -> {n_pairs} pairs")

total_duration = n_pairs * HOLD_TIME
total_frames = int(FPS * total_duration)
print(f"Duration: {total_duration:.1f}s  |  Frames: {total_frames}")

print("Loading brain mask surface...")
brain_mesh = nifti_to_surface(args.brain_mask)
if brain_mesh is not None:
    brain_mesh = brain_mesh.smooth(n_iter=50)

print("Loading bundle surfaces...")
bundle_meshes = []
for i, f in enumerate(bundle_files):
    name = os.path.basename(f)
    print(f"  [{i + 1:3d}/{n_bundles}] {name}", end="")
    mesh = nifti_to_surface(f)
    if mesh is not None:
        mesh = mesh.smooth(n_iter=20)
        print(f"  ({mesh.n_points} pts)")
    else:
        print("  (empty — skipped)")
    bundle_meshes.append(mesh)

plotter = pv.Plotter(off_screen=True, window_size=RESOLUTION)
plotter.set_background(BACKGROUND_COLOR)

# Glass brain
if brain_mesh is not None:
    plotter.add_mesh(
        brain_mesh, opacity=0.08, color="white",
        smooth_shading=True, ambient=0.3, diffuse=0.6,
    )

# Bundle actors — all start invisible (opacity 0)
actors = []
for i, mesh in enumerate(bundle_meshes):
    if mesh is None:
        actors.append(None)
        continue
    pair_idx = i // 2
    color = BLUE if pair_idx % 2 == 0 else ORANGE
    actor = plotter.add_mesh(
        mesh, opacity=0.0, color=color,
        smooth_shading=True,
        specular=0.5, specular_power=20,
        ambient=0.3, diffuse=0.7,
    )
    actors.append(actor)

plotter.reset_camera()
plotter.camera.elevation = 0
plotter.camera.azimuth = 90
plotter.camera.Zoom(1.2)

print("Rendering...")
frames = []
prev_pair = -1

for fi in range(total_frames):
    t = fi / FPS

    # Continuous rotation
    plotter.camera_position = "xz"
    plotter.camera.azimuth = ROTATION_SPEED * t

    # Hard cut: which pair is on right now?
    current_pair = min(int(t / HOLD_TIME), n_pairs - 1)

    if current_pair != prev_pair:
        # Hide previous pair
        if prev_pair >= 0:
            for idx in (prev_pair * 2, prev_pair * 2 + 1):
                if 0 <= idx < len(actors) and actors[idx] is not None:
                    actors[idx].GetProperty().SetOpacity(0.0)
        # Show current pair
        for idx in (current_pair * 2, current_pair * 2 + 1):
            if 0 <= idx < len(actors) and actors[idx] is not None:
                actors[idx].GetProperty().SetOpacity(1.0)
        prev_pair = current_pair

    plotter.render()
    frames.append(plotter.screenshot(return_img=True))

    if fi % FPS == 0:
        print(f"  {fi}/{total_frames}  ({100 * fi / total_frames:.0f}%)")

print(f"Saving to {args.output}...")
imageio.mimwrite(args.output, frames, fps=FPS, quality=10)
print("Done!")