replaced nan masking with coord mask + morph opening-closing

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(

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)