From eae34c0b7c478119ea5ab4dac664a8b33d63dcfd Mon Sep 17 00:00:00 2001 From: harmening Date: Fri, 12 Apr 2024 07:56:55 +0200 Subject: [PATCH] add io functions for head models, Adot, einstein optodes --- src/cedalion/imagereco/forward_model.py | 67 ++++++++++++++++++++++++- src/cedalion/io/__init__.py | 2 + src/cedalion/io/forward_model.py | 18 +++++++ src/cedalion/io/photogrammetry.py | 44 ++++++++++++++++ 4 files changed, 130 insertions(+), 1 deletion(-) create mode 100644 src/cedalion/io/forward_model.py create mode 100644 src/cedalion/io/photogrammetry.py diff --git a/src/cedalion/imagereco/forward_model.py b/src/cedalion/imagereco/forward_model.py index 4c86bcc..1d13263 100644 --- a/src/cedalion/imagereco/forward_model.py +++ b/src/cedalion/imagereco/forward_model.py @@ -7,6 +7,7 @@ import pandas as pd import pmcx import scipy.sparse +import trimesh import xarray as xr import cedalion @@ -27,7 +28,7 @@ class TwoSurfaceHeadModel: segmentation_masks: xr.DataArray brain: cdc.Surface scalp: cdc.Surface - landmarks: Optional[cdt.LabeledPointCloud] + landmarks: cdt.LabeledPointCloud t_ijk2ras: cdt.AffineTransform t_ras2ijk: cdt.AffineTransform voxel_to_vertex_brain: scipy.sparse.spmatrix @@ -150,6 +151,70 @@ def apply_transform(self, transform: cdt.AffineTransform) -> "TwoSurfaceHeadMode voxel_to_vertex_brain=self.voxel_to_vertex_brain, voxel_to_vertex_scalp=self.voxel_to_vertex_scalp, ) + + def save(self, foldername: str): + # Add foldername if not + if ((not os.path.exists(foldername)) or \ + (not os.path.isdir(foldername))): + os.mkdir(foldername) + + self.segmentation_masks.to_netcdf(os.path.join(foldername, + "segmentation_masks.nc")) + self.brain.mesh.export(os.path.join(foldername, "brain.ply"), + file_type="ply") + self.scalp.mesh.export(os.path.join(foldername, "scalp.ply"), + file_type="ply") + if self.landmarks is not None: + self.landmarks.to_netcdf(os.path.join(foldername, "landmarks.nc")) + self.t_ijk2ras.to_netcdf(os.path.join(foldername, "t_ijk2ras.nc")) + self.t_ras2ijk.to_netcdf(os.path.join(foldername, "t_ras2ijk.nc")) + scipy.sparse.save_npz(os.path.join(foldername, "voxel_to_vertex_brain.npz"), + self.voxel_to_vertex_brain) + scipy.sparse.save_npz(os.path.join(foldername, "voxel_to_vertex_scalp.npz"), + self.voxel_to_vertex_scalp) + return + + @classmethod + def load(cls, foldername: str): + for fn in ["segmentation_masks.nc", "brain.ply", "scalp.ply", + "t_ijk2ras.nc", "t_ras2ijk.nc", "voxel_to_vertex_brain.npz", + "voxel_to_vertex_scalp.npz"]: + if not os.path.exists(os.path.join(foldername, fn)): + raise ValueError("%s does not exist." % os.path.join(foldername, fn)) + + # Load data from folder + segmentation_masks = xr.open_dataset(os.path.join(foldername, + 'segmentation_masks.nc')) + brain = trimesh.load(os.path.join(foldername, 'brain.ply'), process=False) + scalp = trimesh.load(os.path.join(foldername, 'scalp.ply'), process=False) + if os.path.exists(os.path.join(foldername, 'landmarks.nc')): + landmarks_ijk = xr.open_dataset(os.path.join(foldername, 'landmarks.nc')) + else: + landmarks_ijk = None + t_ijk2ras = xr.open_dataset(os.path.join(foldername, 't_ijk2ras.nc')) + t_ras2ijk = xr.open_dataset(os.path.join(foldername, 't_ras2ijk.nc')) + voxel_to_vertex_brain = scipy.sparse.load_npz(os.path.join(foldername, + 'voxel_to_vertex_brain.npz')) + voxel_to_vertex_scalp = scipy.sparse.load_npz(os.path.join(foldername, + 'voxel_to_vertex_scalp.npz')) + + # Construct TwoSurfaceHeadModel + brain_ijk = cdc.TrimeshSurface(brain, 'ijk', cedalion.units.Unit("1")) + scalp_ijk = cdc.TrimeshSurface(scalp, 'ijk', cedalion.units.Unit("1")) + t_ijk2ras = cdc.affine_transform_from_numpy(np.array(t_ijk2ras.to_dataarray()[0]), "ijk", "unknown", "1", "mm") + t_ras2ijk = xrutils.pinv(t_ijk2ras) + + return cls( + segmentation_masks=segmentation_masks, + brain=brain_ijk, + scalp=scalp_ijk, + landmarks=landmarks_ijk, + t_ijk2ras=t_ijk2ras, + t_ras2ijk=t_ras2ijk, + voxel_to_vertex_brain=voxel_to_vertex_brain, + voxel_to_vertex_scalp=voxel_to_vertex_scalp, + ) + # FIXME maybe this should not be in this class, especially since the # algorithm is not good. diff --git a/src/cedalion/io/__init__.py b/src/cedalion/io/__init__.py index 297fa2d..5103d2d 100644 --- a/src/cedalion/io/__init__.py +++ b/src/cedalion/io/__init__.py @@ -1,3 +1,5 @@ from .snirf import read_snirf from .probe_geometry import read_mrk_json, read_digpts, read_einstar_obj from .anatomy import read_segmentation_masks +from .photogrammetry import read_photogrammetry_einstar, read_einstar, opt_fid_to_xr +from .forward_model import save_Adot, load_Adot diff --git a/src/cedalion/io/forward_model.py b/src/cedalion/io/forward_model.py new file mode 100644 index 0000000..2ebff54 --- /dev/null +++ b/src/cedalion/io/forward_model.py @@ -0,0 +1,18 @@ +import xarray as xr +import cedalion + +def save_Adot(fn: str, Adot: xr.DataArray): + Adot.to_netcdf(fn) + return + +def load_Adot(fn: str): + Adot = xr.open_dataset(fn) + Adot = xr.DataArray( + Adot.to_array()[0], + dims=["channel", "vertex", "wavelength"], + coords={"channel": ("channel", Adot.channel.values), + "wavelength": ("wavelength", Adot.wavelength.values), + "is_brain": ("vertex", Adot.is_brain.values) + } + ) + return Adot diff --git a/src/cedalion/io/photogrammetry.py b/src/cedalion/io/photogrammetry.py new file mode 100644 index 0000000..a3dfefe --- /dev/null +++ b/src/cedalion/io/photogrammetry.py @@ -0,0 +1,44 @@ +import cedalion +import cedalion.dataclasses as cdc +import numpy as np +from collections import OrderedDict + +def read_photogrammetry_einstar(fn): + fiducials, optodes = read_einstar(fn) + fiducials, optodes = opt_fid_to_xr(fiducials, optodes) + return fiducials, optodes + + +def read_einstar(fn): + with open(fn, 'r') as f: + lines = [[l.strip() for l in line.split(',')] for line in f.readlines()] + lines = [[line[0], [float(l) for l in line[1:]]] for line in lines] + assert lines[0][0] == 'Nz' + assert lines[1][0] == 'Iz' + assert lines[2][0] == 'Rpa' + assert lines[3][0] == 'Lpa' + assert lines[4][0] == 'Cz' + fiducials = OrderedDict(lines[:5]) + optodes = OrderedDict(lines[5:]) + return fiducials, optodes + + +def opt_fid_to_xr(fiducials, optodes): + # FIXME: this should get a different CRS + CRS = "ijk" + fiducials = cdc.build_labeled_points(np.array(list(fiducials.values())), + labels=list(fiducials.keys()), + crs=CRS)#, units="mm") + opt = np.array(list(optodes.values())) + types = [] + for lab in list(optodes.keys()): + if lab.startswith('S'): + types.append(cdc.PointType(1)) + elif lab.startswith('D'): + types.append(cdc.PointType(2)) + else: + types.append(cdc.PointType(0)) + optodes = cdc.build_labeled_points(opt, + labels=list(optodes.keys()), + crs=CRS, types=types) + return fiducials, optodes