From e7cc3566f472ca79b58b7fe88f8185f28f73a9cf Mon Sep 17 00:00:00 2001 From: Jordan DeKraker Date: Thu, 23 Jan 2025 16:25:58 -0500 Subject: [PATCH] replaced nan masking with coord mask + morph opening-closing --- hippunfold/workflow/rules/native_surf.smk | 1 + hippunfold/workflow/scripts/gen_isosurface.py | 177 ++++++++++-------- 2 files changed, 99 insertions(+), 79 deletions(-) diff --git a/hippunfold/workflow/rules/native_surf.smk b/hippunfold/workflow/rules/native_surf.smk index c7c6271a..62dae2fa 100644 --- a/hippunfold/workflow/rules/native_surf.smk +++ b/hippunfold/workflow/rules/native_surf.smk @@ -204,6 +204,7 @@ rule gen_native_mesh: params: threshold=lambda wildcards: surf_thresholds[wildcards.surfname], decimate_percent=0, # not enabled + morph_openclose_dist=2, # mm output: surf_gii=temp( bids( diff --git a/hippunfold/workflow/scripts/gen_isosurface.py b/hippunfold/workflow/scripts/gen_isosurface.py index 0ee9ff6c..9ad01278 100644 --- a/hippunfold/workflow/scripts/gen_isosurface.py +++ b/hippunfold/workflow/scripts/gen_isosurface.py @@ -1,14 +1,10 @@ import pyvista as pv import nibabel as nib import numpy as np -from vtk import vtkNIFTIImageReader +from copy import deepcopy -def write_surface_to_gifti(in_surface, out_surf_gii): - - faces = in_surface.faces - faces = faces.reshape((int(faces.shape[0] / 4), 4))[:, 1:4] - points = in_surface.points +def write_surface_to_gifti(points, faces, out_surf_gii): points_darray = nib.gifti.GiftiDataArray( data=points, intent="NIFTI_INTENT_POINTSET", datatype="NIFTI_TYPE_FLOAT32" @@ -25,66 +21,86 @@ def write_surface_to_gifti(in_surface, out_surf_gii): gifti.to_filename(out_surf_gii) -def remove_nan_points_faces(vertices, faces): - # Step 1: Identify valid vertices (no NaN values) - nan_mask = np.isnan(vertices).any(axis=1) - valid_vertices = ~nan_mask # True for valid rows - valid_indices = np.where(valid_vertices)[0] +def remove_nan_vertices(vertices, faces): + """ + Removes vertices containing NaNs and updates faces accordingly. + + Parameters: + - vertices (np.ndarray): (N, 3) array of vertex positions. + - faces (np.ndarray): (M, 3) array of triangular face indices. - # Step 2: Create a mapping from old to new indices - new_indices_map = -np.ones( - vertices.shape[0], dtype=int - ) # Default to -1 for invalid vertices - new_indices_map[valid_indices] = np.arange(len(valid_indices)) + Returns: + - new_vertices (np.ndarray): Filtered (N', 3) array of valid vertex positions. + - new_faces (np.ndarray): Filtered (M', 3) array of updated face indices. + """ + # Identify valid (non-NaN) vertices + valid_mask = ~np.isnan(vertices).any(axis=1) - # Step 3: Update the faces array to remove references to invalid vertices - # Replace old indices with new ones, and remove faces with invalid vertices - new_faces = [] - for face in faces: - # Map old indices to new ones - mapped_face = new_indices_map[face] - if np.all(mapped_face >= 0): # Include only faces with all valid vertices - new_faces.append(mapped_face) + # Create a mapping from old indices to new indices + new_indices = np.full( + vertices.shape[0], -1, dtype=int + ) # Default -1 for invalid ones + new_indices[valid_mask] = np.arange(valid_mask.sum()) # Renumber valid vertices - new_faces = np.array(new_faces) + # Filter out faces that reference removed vertices + valid_faces_mask = np.all( + valid_mask[faces], axis=1 + ) # Keep only faces with valid vertices + new_faces = new_indices[faces[valid_faces_mask]] # Remap face indices - # Step 4: Remove invalid vertices from the array - new_vertices = vertices[valid_vertices] + # Filter vertices + new_vertices = vertices[valid_mask] - return (new_vertices, new_faces) + return new_vertices, new_faces -from scipy.ndimage import binary_dilation +import scipy.sparse as sp +from scipy.sparse.csgraph import dijkstra -def get_adjacent_voxels(mask_a, mask_b): +def compute_geodesic_distances(vertices, faces, source_indices): """ - Create a mask for voxels where label A is adjacent to label B. + Computes geodesic distances from a set of source vertices to all other vertices on a 3D surface mesh. Parameters: - - mask_a (np.ndarray): A 3D binary mask for label A. - - mask_b (np.ndarray): A 3D binary mask for label B. + - vertices (np.ndarray): (N, 3) array of vertex positions. + - faces (np.ndarray): (M, 3) array of triangular face indices. + - source_indices (list or np.ndarray): Indices of source vertices. Returns: - - np.ndarray: A 3D mask where adjacent voxels for label A and label B are marked as True. + - distances (np.ndarray): (N,) array of geodesic distances from the source vertices. """ - # Dilate each mask to identify neighboring regions - dilated_a = binary_dilation(mask_a) - dilated_b = binary_dilation(mask_b) + num_vertices = len(vertices) - # Find adjacency: voxels of A touching B and B touching A - adjacency_mask = (dilated_a.astype("bool") & mask_b.astype("bool")) | ( - dilated_b.astype("bool") & mask_a.astype("bool") - ) + # Create adjacency matrix + row, col, weight = [], [], [] + for f in faces: + for i in range(3): + v1, v2 = f[i], f[(i + 1) % 3] # Pairwise edges in the triangle + dist = np.linalg.norm(vertices[v1] - vertices[v2]) # Euclidean edge length + row.append(v1) + col.append(v2) + weight.append(dist) + row.append(v2) + col.append(v1) + weight.append(dist) # Ensure symmetry + + graph = sp.csr_matrix((weight, (row, col)), shape=(num_vertices, num_vertices)) + + # Compute geodesic distances using Dijkstra's algorithm + distances = dijkstra(csgraph=graph, directed=False, indices=source_indices) - return adjacency_mask + # If multiple sources, take the minimum distance to any of them + if isinstance(source_indices, (list, np.ndarray)) and len(source_indices) > 1: + distances = np.min(distances, axis=0) + + return distances # Load the coords image coords_img = nib.load(snakemake.input.coords) coords = coords_img.get_fdata() - # Load the nan mask nan_mask_img = nib.load(snakemake.input.nan_mask) nan_mask = nan_mask_img.get_fdata() @@ -111,12 +127,8 @@ def get_adjacent_voxels(mask_a, mask_b): # update the coords data to add the nans and sink coords[nan_mask == 1] = np.nan -coords[sink_mask == 1] = 1.1 # since sink being zero creates a false boundary - -# we also need to use a nan mask for the voxels where src and sink meet directly -# (since this is another false boundary).. -src_sink_nan_mask = get_adjacent_voxels(sink_mask, src_mask) -coords[src_sink_nan_mask == 1] = np.nan +coords[src_mask == 1] = -0.1 +coords[sink_mask == 1] = 1.1 # Add the scalar field @@ -130,41 +142,48 @@ def get_adjacent_voxels(mask_a, mask_b): # the contour function produces the isosurface - -surface = tfm_grid.contour([snakemake.params.threshold], method="contour").decimate( - 0.9 -) # fill_holes(snakemake.params.max_hole_size) -# surface = tfm_grid.contour([snakemake.params.threshold],method='contour').fill_holes(snakemake.params.max_hole_size) - -# surface = surface.decimate(float(snakemake.params.decimate_percent) / 100.0) - +surface = tfm_grid.contour([snakemake.params.threshold], method="contour").decimate(0.9) # faces from pyvista surface are formatted with number of verts each row # reshape and remove the first col to get Nx3 faces = surface.faces faces = faces.reshape((int(faces.shape[0] / 4), 4))[:, 1:4] - points = surface.points +points, faces = remove_nan_vertices(points, faces) + +## JD clean - instead of trimming surfaces with a nan mask, we +# keep vertices that overlap with good coord values. We then apply +# some surface-based morphological opening and closing to keep +# vertices along holes in the dg + +# this is equivalent to wb_command -volume-to-surface-mapping -enclosing +# apply inverse affine to surface to get back to matrix space +V = deepcopy(points) +V[:, :] = V - affine[:3, 3].T +for xyz in range(3): + V[:, xyz] = V[:, xyz] * (1 / affine[xyz, xyz]) +V = V.astype(int) +# sample coords +coord_at_V = np.zeros((len(V))) +for i in range(len(V)): + coord_at_V[i] = coords[ + V[i, 0], V[i, 1], V[i, 2] + ] # really hope there's no x-y switching fuckery here! + +# keep vertices that are in a nice coordinate range +good_v = np.where(np.logical_and(coord_at_V < 0.9, coord_at_V > 0.1))[0] + +# morphological open +maxdist = compute_geodesic_distances(points, faces, good_v) +bad_v = np.where(maxdist > snakemake.params.morph_openclose_dist)[0] + +# morphological close +maxdist = compute_geodesic_distances(points, faces, bad_v) +bad_v = np.where(maxdist < snakemake.params.morph_openclose_dist)[0] + +# toss bad vertices +points[bad_v, :] = np.nan +points, faces = remove_nan_vertices(points, faces) -# with nans in background we end up with nan vertices, we can remove -# these to end up with an open contour.. -new_points, new_faces = remove_nan_points_faces(points, faces) - - -# Step 1: Prepare the PolyData -# PyVista expects faces in a flat array with the number of points in each face as the first value -faces_flat = np.hstack( - [[3] + list(face) for face in new_faces] -) # Add '3' for triangular faces - -# Create a new PolyData object -polydata = pv.PolyData(new_points, faces_flat) - -# Ensure the PolyData is clean (optional) -# polydata.clean(inplace=True) # Removes unused points, degenerate cells, etc. - -# Apply decimation (optional) -# polydata.decimate_pro(snakemake.params.decimate_percent / 100.0, inplace=True) - # write to gifti -write_surface_to_gifti(polydata, snakemake.output.surf_gii) +write_surface_to_gifti(points, faces, snakemake.output.surf_gii)