diff --git a/environment_dev.yml b/environment_dev.yml index c2d64425..5e8ac284 100644 --- a/environment_dev.yml +++ b/environment_dev.yml @@ -54,6 +54,6 @@ dependencies: - pywavefront==1.3 - setuptools-scm - snirf==0.8 - - pmcx==0.3.3 + - pmcx==0.3.4 + - pmcxcl==0.2.0 - sphinxcontrib-bibtex - diff --git a/environment_doc.yml b/environment_doc.yml index 2cb5dbd7..9d3e129a 100644 --- a/environment_doc.yml +++ b/environment_doc.yml @@ -55,4 +55,5 @@ dependencies: - setuptools-scm==7.1.0 - snirf==0.7.4 - pmcx + - pmcxcl - open3d==0.16 diff --git a/src/cedalion/dataclasses/__init__.py b/src/cedalion/dataclasses/__init__.py index d5863df1..5005c2fd 100644 --- a/src/cedalion/dataclasses/__init__.py +++ b/src/cedalion/dataclasses/__init__.py @@ -9,6 +9,7 @@ TrimeshSurface, VTKSurface, affine_transform_from_numpy, + Voxels, ) from .schemas import ( build_labeled_points, diff --git a/src/cedalion/dataclasses/geometry.py b/src/cedalion/dataclasses/geometry.py index 7728b9bc..d013827e 100644 --- a/src/cedalion/dataclasses/geometry.py +++ b/src/cedalion/dataclasses/geometry.py @@ -96,6 +96,63 @@ def snap(self, points: cdt.LabeledPointCloud): return snapped +@dataclass +class Voxels(): + """3D voxels represented by a np.array. + + Attributes: + voxels (np.ndarray): The voxels. + crs (str): The coordinate reference system of the voxels. + units (pint.Unit): The units of the voxels. + """ + voxels: np.ndarray + crs: str + units: pint.Unit + + @property + def vertices(self) -> cdt.LabeledPointCloud: + result = xr.DataArray( + self.voxels, + dims=["label", self.crs], + coords={"label": np.arange(len(self.voxels))}, + attrs={"units": self.units}, + ) + result = result.pint.quantify() + + return result + + @property + def nvertices(self) -> int: + return len(self.voxels) + + def apply_transform(self, transform: cdt.AffineTransform) -> "Voxels": + # convert to homogeneous coordinates + num, dim = self.voxels.shape + hom = np.ones((num,dim+1)) + hom[:,:3] = self.voxels + # apply transformation + hom = (transform.pint.dequantify().values.dot(hom.T)).T + # backtransformation + transformed = np.array([hom[i,:3] / hom[i,3] for i in range(hom.shape[0])]) + + new_units = self.units * transform.pint.units + new_crs = transform.dims[0] + + return Voxels(transformed, new_crs, new_units) + + def _build_kdtree(self): + self._kdtree = KDTree(self.voxels) + + def __post_init__(self): + self._kdtree = None + + @property + def kdtree(self): + if self._kdtree is None: + self._build_kdtree() + return self._kdtree + + @dataclass class TrimeshSurface(Surface): """A surface represented by a trimesh object. diff --git a/src/cedalion/geometry/segmentation.py b/src/cedalion/geometry/segmentation.py index ef91962b..b7152c64 100644 --- a/src/cedalion/geometry/segmentation.py +++ b/src/cedalion/geometry/segmentation.py @@ -12,6 +12,42 @@ import cedalion.dataclasses as cdc +def voxels_from_segmentation( + segmentation_mask: xr.DataArray, + segmentation_types: List[str], + isovalue=0.9, + fill_holes_in_mask=False, +) -> cdc.Voxels: + """Generate voxels from a segmentation mask. + + Args: + segmentation_mask : xr.DataArray + Segmentation mask. + segmentation_types : List[str] + List of segmentation types. + isovalue : float, optional + Isovalue for marching cubes, by default 0.9. + fill_holes_in_mask : bool, optional + Fill holes in the mask, by default False. + + Returns: + cdc.Voxels + Voxels in voxel space. + """ + combined_mask = ( + segmentation_mask.sel(segmentation_type=segmentation_types) + .any("segmentation_type") + .values + ) + + if fill_holes_in_mask: + combined_mask = binary_fill_holes(combined_mask).astype(combined_mask.dtype) + + voxels = np.argwhere(combined_mask) + + return cdc.Voxels(voxels, "ijk", cedalion.units.Unit("1")) + + def surface_from_segmentation( segmentation_mask: xr.DataArray, segmentation_types: List[str], @@ -32,6 +68,7 @@ def surface_from_segmentation( Returns: A cedalion.Surface object. """ + combined_mask = ( segmentation_mask.sel(segmentation_type=segmentation_types) .any("segmentation_type") @@ -59,6 +96,7 @@ def cell_coordinates(volume, flat: bool = False): Returns: xr.DataArray: A DataArray with the coordinates of the cells in the volume. """ + # coordinates in voxel space i = np.arange(volume.shape[0]) j = np.arange(volume.shape[1]) diff --git a/src/cedalion/imagereco/forward_model.py b/src/cedalion/imagereco/forward_model.py index 006417d3..b331fd3f 100644 --- a/src/cedalion/imagereco/forward_model.py +++ b/src/cedalion/imagereco/forward_model.py @@ -647,7 +647,7 @@ def _get_unitinmm(self): length = xrutils.norm(pts_ras[1] - pts_ras[0], pts_ras.points.crs) return length.pint.magnitude.item() - def _get_fluence_from_mcx(self, i_optode: int, nphoton: int): + def _get_fluence_from_mcx(self, i_optode: int, **kwargs): """Run MCX simulation to get fluence for one optode. Args: @@ -658,8 +658,10 @@ def _get_fluence_from_mcx(self, i_optode: int, nphoton: int): np.ndarray: Fluence in each voxel. """ + kwargs.setdefault("nphoton", 1e8) + cfg = { - "nphoton": nphoton, + "nphoton": kwargs['nphoton'], "vol": self.volume, "tstart": 0, "tend": 5e-9, @@ -670,16 +672,24 @@ def _get_fluence_from_mcx(self, i_optode: int, nphoton: int): "issrcfrom0": 1, "isnormalized": 1, "outputtype": "fluence", - "seed": int(np.floor(np.random.rand() * 1e7)), - "issavedet": 1, + "issavedet": 0, "unitinmm": self.unitinmm, } - import pmcx - result = pmcx.run(cfg) + # merging default cfg with additional positional arguments + + cfg = { **cfg, **kwargs } + + # if pmcx fails, try pmcxcl + + if "cuda" in cfg and cfg["cuda"]: + import pmcx + result = pmcx.run(cfg) + else: + import pmcxcl + result = pmcxcl.run(cfg) fluence = result["flux"][:, :, :, 0] # there is only one time bin - fluence = fluence * cfg["tstep"] / result["stat"]["normalizer"] return fluence @@ -722,11 +732,13 @@ def _fluence_at_optodes(self, fluence, emitting_opt): return result - def compute_fluence_mcx(self, nphoton: int = 1e8): + def compute_fluence_mcx(self, **kwargs): """Compute fluence for each channel and wavelength using MCX package. Args: nphoton (int): Number of photons to simulate. + along with other pmcx/pmcxcl accepted input fields, + see https://pypi.org/project/pmcx/ Returns: xr.DataArray: Fluence in each voxel for each channel and wavelength. @@ -763,9 +775,9 @@ def compute_fluence_mcx(self, nphoton: int = 1e8): label = self.optode_pos.label.values[i_opt] print(f"simulating fluence for {label}. {i_opt+1} / {n_optodes}") - # run MCX + # run MCX or MCXCL # shape: [i,j,k] - fluence = self._get_fluence_from_mcx(i_opt, nphoton=nphoton) + fluence = self._get_fluence_from_mcx(i_opt, **kwargs) # FIXME shortcut: currently tissue props are wavelength independent -> copy for i_wl in range(n_wavelength): diff --git a/src/cedalion/io/forward_model.py b/src/cedalion/io/forward_model.py index f33fa578..408dae91 100644 --- a/src/cedalion/io/forward_model.py +++ b/src/cedalion/io/forward_model.py @@ -44,7 +44,7 @@ def load_Adot(fn: str): def save_fluence(fn : str, fluence_all, fluence_at_optodes): """Save forward model computation results. - This method uses a lossy compressions algorithm to reduce file size. + This method uses a lossy compression algorithm to reduce file size. """ with h5py.File(fn, "w") as f: