add io functions for head models, Adot, einstein optodes

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.

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

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

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