From e6d8ae417f34306ff425a0d1cdac817cc11f9f57 Mon Sep 17 00:00:00 2001 From: harmening Date: Tue, 4 Nov 2025 23:25:45 +0100 Subject: [PATCH 1/2] enable 2SHM to use voxel parcellations for vox2surf mapping --- src/cedalion/imagereco/forward_model.py | 33 +++++++++++++++----- src/cedalion/imagereco/utils.py | 41 +++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 7 deletions(-) diff --git a/src/cedalion/imagereco/forward_model.py b/src/cedalion/imagereco/forward_model.py index 7893e605..be2af2bb 100644 --- a/src/cedalion/imagereco/forward_model.py +++ b/src/cedalion/imagereco/forward_model.py @@ -243,6 +243,7 @@ def from_surfaces( scalp_face_count: int | None = 60000, fill_holes: bool = False, parcel_file: Path | str | None = None, + parcel_volume_file: Path | str | None = None, ) -> "TwoSurfaceHeadModel": """Constructor from seg.masks, brain and head surfaces as gained from MRI scans. @@ -263,6 +264,7 @@ def from_surfaces( scalp_face_count (Optional[int]): Number of faces for the scalp surface. fill_holes (bool): Whether to fill holes in the segmentation masks. parcel_file: Path to parcel json file. + parcel_volume_file: Path to parcel nifiti file (annotated voxels). Returns: TwoSurfaceHeadModel: An instance of the TwoSurfaceHeadModel class. @@ -343,19 +345,36 @@ def from_surfaces( "segmentation_type" ) - voxel_to_vertex_brain = map_segmentation_mask_to_surface( - brain_mask, t_ijk2ras, brain_ijk.apply_transform(t_ijk2ras) - ) - voxel_to_vertex_scalp = map_segmentation_mask_to_surface( - scalp_mask, t_ijk2ras, scalp_ijk.apply_transform(t_ijk2ras) - ) - # load parcellations if parcel_file is not None: parcels = cedalion.io.read_parcellations(parcel_file) assert len(parcels) == brain_ijk.nvertices brain_ijk.vertex_coords["parcel"] = np.asarray(parcels.Label.tolist()) + if parcel_volume_file is not None: + import nibabel as nib + voxel_parcels = nib.load(parcel_volume_file) + affine = voxel_parcels.affine + voxel_parcels = voxel_parcels.get_fdata() + labels = np.unique(voxel_parcels.astype(int)) + if os.path.exists(parcel_volume_file.replace('.nii.gz', '_labels.csv')): + with open(parcel_volume_file.replace('.nii.gz', '_labels.csv'), 'r') as f: + lines = [l.split() for l in f.readlines()] + csv_cbig = {int(l[0]): l[1] for l in lines} + parcels_dict = parcels.Label.to_dict() + for i, l in csv_cbig.items(): + assert csv_cbig[i] == l + assert brain_mask.shape == voxel_parcels.shape + assert (t_ijk2ras.values == affine).all() + + voxel_to_vertex_brain = map_segmentation_mask_to_surface( + brain_mask, t_ijk2ras, brain_ijk.apply_transform(t_ijk2ras), + parcels_vox=voxel_parcels.astype(int), parcels_verts=parcels + ) + voxel_to_vertex_scalp = map_segmentation_mask_to_surface( + scalp_mask, t_ijk2ras, scalp_ijk.apply_transform(t_ijk2ras) + ) + return cls( segmentation_masks=segmentation_masks, diff --git a/src/cedalion/imagereco/utils.py b/src/cedalion/imagereco/utils.py index 4e5529f6..ffbdbc86 100644 --- a/src/cedalion/imagereco/utils.py +++ b/src/cedalion/imagereco/utils.py @@ -18,6 +18,8 @@ def map_segmentation_mask_to_surface( segmentation_mask: xr.DataArray, transform_vox2ras: cdt.AffineTransform, # FIXME surface: cdc.Surface, + parcels_vox: np.ndarray = None, + parcels_verts: xr.DataArray = None, ): """Find for each voxel the closest vertex on the surface. @@ -27,6 +29,10 @@ def map_segmentation_mask_to_surface( transform_vox2ras (xr.DataArray): The affine transformation from voxel to RAS space. surface (cedalion.dataclasses.Surface): The surface to map the voxels to. + parcels_vox (np.ndarray, optional): An array of shape (nx, ny, nz) containing + the parcel label indices for each voxel. + parcels_verts (xr.DataArray, optional): An array of shape (nvertices,) containing + the parcellation information for each brain surface vertex. Returns: coo_array: A sparse matrix of shape (ncells, nvertices) that maps voxels to @@ -51,6 +57,41 @@ def map_segmentation_mask_to_surface( cell_coords.values[cell_indices, :], workers=-1 ) + if parcels_vox is not None and parcels_verts is not None: + # overwrite voxel labels if in segmentation mask not in brain tissue + fs_num_labeled_vox = np.sum(np.flatnonzero(parcels_vox)) + parcels_vox *= segmentation_mask.values + print("Num of labeled voxels before seg-masking: %d\nNum of labeled voxels after seg-masking: %d" % (fs_num_labeled_vox, np.sum(np.flatnonzero(parcels_vox)))) + + # if parcellation is provided, overwrite vertex_indices with mapping + # constraint to vertices-mapping within the same parcel + for parcel_id in np.unique(parcels_vox): + if parcel_id == 0: + continue + # get cell indices within this parcel + parcels_vox_flat = parcels_vox.flatten() + parcel_cell_indices = np.argwhere(parcels_vox_flat[cell_indices] == parcel_id)[:, 0] + if len(parcel_cell_indices) == 0: + continue + # get vertices within this parcel + parcel_vertex_indices = np.where( + parcels_verts.index == parcel_id + )[0] + if len(parcel_vertex_indices) == 0: + continue + # build a KDTree for the parcel vertices + from scipy.spatial import KDTree + parcel_tree = KDTree(surface.vertices[parcel_vertex_indices, :]) + # query the parcel_tree for the parcel_cell_indices + dists_parcel, vertex_indices_parcel = parcel_tree.query( + cell_coords.values[parcel_cell_indices, :], workers=-1 + ) + # map back to global vertex indices + global_vertex_indices_parcel = parcel_vertex_indices[vertex_indices_parcel] + # update vertex_indices for these cell indices + + vertex_indices[parcel_cell_indices] = global_vertex_indices_parcel + # construct a sparse matrix of shape (ncells, nvertices) # that maps voxels to cells map_voxel_to_vertex = coo_array( From a858242dce3c1595dd8ccaa888410910e18c57cd Mon Sep 17 00:00:00 2001 From: harmening Date: Fri, 14 Nov 2025 11:50:39 +0100 Subject: [PATCH 2/2] bypass hack if parcel labels are correctly provided --- src/cedalion/dot/head_model.py | 7 ++++++- src/cedalion/io/anatomy.py | 3 ++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/cedalion/dot/head_model.py b/src/cedalion/dot/head_model.py index 9e0c3393..d9b8e69b 100644 --- a/src/cedalion/dot/head_model.py +++ b/src/cedalion/dot/head_model.py @@ -334,6 +334,8 @@ def from_surfaces( parcels = cedalion.io.read_parcellations(parcel_file) assert len(parcels) == brain_ijk.nvertices brain_ijk.vertex_coords["parcel"] = np.asarray(parcels.Label.tolist()) + else: + parcels = None if parcel_volume_file is not None: import nibabel as nib @@ -350,10 +352,13 @@ def from_surfaces( assert csv_cbig[i] == l assert brain_mask.shape == voxel_parcels.shape assert (t_ijk2ras.values == affine).all() + voxel_parcels = voxel_parcels.astype(int) + else: + voxel_parcels = None voxel_to_vertex_brain = map_segmentation_mask_to_surface( brain_mask, t_ijk2ras, brain_ijk.apply_transform(t_ijk2ras), - parcels_vox=voxel_parcels.astype(int), parcels_verts=parcels + parcels_vox=voxel_parcels, parcels_verts=parcels ) voxel_to_vertex_scalp = map_segmentation_mask_to_surface( scalp_mask, t_ijk2ras, scalp_ijk.apply_transform(t_ijk2ras) diff --git a/src/cedalion/io/anatomy.py b/src/cedalion/io/anatomy.py index fe301962..5744d566 100644 --- a/src/cedalion/io/anatomy.py +++ b/src/cedalion/io/anatomy.py @@ -168,6 +168,7 @@ def read_parcellations(parcel_file: str | Path) -> pd.DataFrame: parcels["Vertices"] = parcels["Vertices"].astype(int) parcels = parcels.sort_values("Vertices") - parcels["Label"] = parcels["Label"].apply(lambda x: "_".join(x.split(" ")) + "H") + if not parcels["Label"].values[1].endswith('H'): + parcels["Label"] = parcels["Label"].apply(lambda x: "_".join(x.split(" ")) + "H") return parcels