junifer 0.0.3.dev186__py3-none-any.whl → 0.0.4__py3-none-any.whl

junifer/data/masks.py

@@ -1,8 +1,10 @@
 """Provide functions for masks."""
 
 # Authors: Federico Raimondo <f.raimondo@fz-juelich.de>
+#          Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
+import typing
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
@@ -18,29 +20,111 @@ from typing import (
 import nibabel as nib
 import numpy as np
 from nilearn.datasets import fetch_icbm152_brain_gm_mask
-from nilearn.image import resample_to_img
+from nilearn.image import get_data, new_img_like, resample_to_img
 from nilearn.masking import (
     compute_background_mask,
-    compute_brain_mask,
     compute_epi_mask,
     intersect_masks,
 )
 
-from ..utils.logging import logger, raise_error
+from ..pipeline import WorkDirManager
+from ..utils import logger, raise_error, run_ext_cmd, warn_with_log
+from .template_spaces import get_template, get_xfm
 from .utils import closest_resolution
 
 
 if TYPE_CHECKING:
     from nibabel import Nifti1Image
 
-# Path to the VOIs
+# Path to the masks
 _masks_path = Path(__file__).parent / "masks"
 
 
+def compute_brain_mask(
+    target_data: Dict[str, Any],
+    extra_input: Optional[Dict[str, Any]] = None,
+    mask_type: str = "brain",
+    threshold: float = 0.5,
+) -> "Nifti1Image":
+    """Compute the whole-brain, grey-matter or white-matter mask.
+
+    This mask is calculated using the template space and resolution as found
+    in the ``target_data``.
+
+    Parameters
+    ----------
+    target_data : dict
+        The corresponding item of the data object for which mask will be
+        loaded.
+    extra_input : dict, optional
+        The other fields in the data object. Useful for accessing other data
+        types (default None).
+    mask_type : {"brain", "gm", "wm"}, optional
+        Type of mask to be computed:
+
+        * "brain" : whole-brain mask
+        * "gm" : grey-matter mask
+        * "wm" : white-matter mask
+
+        (default "brain").
+    threshold : float, optional
+        The value under which the template is cut off (default 0.5).
+
+    Returns
+    -------
+    Nifti1Image
+        The mask (3D image).
+
+    Raises
+    ------
+    ValueError
+        If ``mask_type`` is invalid or
+        if ``extra_input`` is None when ``target_data``'s space is native.
+
+    """
+    logger.debug(f"Computing {mask_type} mask")
+
+    if mask_type not in ["brain", "gm", "wm"]:
+        raise_error(f"Unknown mask type: {mask_type}")
+
+    # Check pre-requirements for space manipulation
+    target_space = target_data["space"]
+    # Set target standard space to target space
+    target_std_space = target_space
+    # Extra data type requirement check if target space is native
+    if target_space == "native":
+        # Check for extra inputs
+        if extra_input is None:
+            raise_error(
+                "No extra input provided, requires `Warp` "
+                "data type to infer target template space."
+            )
+        # Set target standard space to warp file space source
+        target_std_space = extra_input["Warp"]["src"]
+
+    # Fetch template in closest resolution
+    template = get_template(
+        space=target_std_space,
+        target_data=target_data,
+        extra_input=extra_input,
+        template_type=mask_type if mask_type in ["gm", "wm"] else "T1w",
+    )
+    # Resample template to target image
+    target_img = target_data["data"]
+    resampled_template = resample_to_img(
+        source_img=template, target_img=target_img
+    )
+
+    # Threshold and get mask
+    mask = (get_data(resampled_template) >= threshold).astype("int8")
+
+    return new_img_like(target_img, mask)  # type: ignore
+
+
 def _fetch_icbm152_brain_gm_mask(
     target_img: "Nifti1Image",
     **kwargs,
-):
+) -> "Nifti1Image":
     """Fetch ICBM152 brain mask and resample.
 
     Parameters
@@ -55,7 +139,21 @@ def _fetch_icbm152_brain_gm_mask(
     -------
     nibabel.Nifti1Image
         The resampled mask.
+
+    Warns
+    -----
+    DeprecationWarning
+        If this function is used.
+
     """
+    warn_with_log(
+        msg=(
+            "It is recommended to use ``compute_brain_mask`` with "
+            "``mask_type='gm'``. This function will be removed in the next "
+            "release. For now, it's available for backward compatibility."
+        ),
+        category=DeprecationWarning,
+    )
     mask = fetch_icbm152_brain_gm_mask(**kwargs)
     mask = resample_to_img(
         mask, target_img, interpolation="nearest", copy=True
@@ -66,29 +164,40 @@ def _fetch_icbm152_brain_gm_mask(
 # A dictionary containing all supported masks and their respective file or
 # data.
 
+# Each entry is a dictionary that must contain at least the following keys:
+# * 'family': the mask's family name (e.g., 'Vickery-Patil', 'Callable')
+# * 'space': the mask's space (e.g., 'MNI', 'inherit')
+
 # The built-in masks are files that are shipped with the package in the
 # data/masks directory. The user can also register their own masks.
 
 # Callable masks should be functions that take at least one parameter:
 # * `target_img`: the image to which the mask will be applied.
 _available_masks: Dict[str, Dict[str, Any]] = {
-    "GM_prob0.2": {"family": "Vickery-Patil"},
-    "GM_prob0.2_cortex": {"family": "Vickery-Patil"},
+    "GM_prob0.2": {"family": "Vickery-Patil", "space": "IXI549Space"},
+    "GM_prob0.2_cortex": {
+        "family": "Vickery-Patil",
+        "space": "IXI549Space",
+    },
     "compute_brain_mask": {
         "family": "Callable",
         "func": compute_brain_mask,
+        "space": "inherit",
     },
     "compute_background_mask": {
         "family": "Callable",
         "func": compute_background_mask,
+        "space": "inherit",
     },
     "compute_epi_mask": {
         "family": "Callable",
         "func": compute_epi_mask,
+        "space": "inherit",
     },
     "fetch_icbm152_brain_gm_mask": {
         "family": "Callable",
         "func": _fetch_icbm152_brain_gm_mask,
+        "space": "MNI152NLin2009aAsym",
     },
 }
 
@@ -96,6 +205,7 @@ _available_masks: Dict[str, Dict[str, Any]] = {
 def register_mask(
     name: str,
     mask_path: Union[str, Path],
+    space: str,
     overwrite: bool = False,
 ) -> None:
     """Register a custom user mask.
@@ -106,6 +216,8 @@ def register_mask(
         The name of the mask.
     mask_path : str or pathlib.Path
         The path to the mask file.
+    space : str
+        The space of the mask, for e.g., "MNI152NLin6Asym".
     overwrite : bool, optional
         If True, overwrite an existing mask with the same name.
         Does not apply to built-in mask (default False).
@@ -115,6 +227,7 @@ def register_mask(
     ValueError
         If the mask name is already registered and overwrite is set to
         False or if the mask name is a built-in mask.
+
     """
     # Check for attempt of overwriting built-in parcellations
     if name in _available_masks:
@@ -122,11 +235,11 @@ def register_mask(
             logger.info(f"Overwriting {name} mask")
             if _available_masks[name]["family"] != "CustomUserMask":
                 raise_error(
-                    f"Cannot overwrite {name} mask. " "It is a built-in mask."
+                    f"Cannot overwrite {name} mask. It is a built-in mask."
                 )
         else:
             raise_error(
-                f"Mask {name} already registered. Set `overwrite=True`"
+                f"Mask {name} already registered. Set `overwrite=True` "
                 "to update its value."
             )
     # Convert str to Path
@@ -136,6 +249,7 @@ def register_mask(
     _available_masks[name] = {
         "path": str(mask_path.absolute()),
         "family": "CustomUserMask",
+        "space": space,
     }
 
 
@@ -146,11 +260,12 @@ def list_masks() -> List[str]:
     -------
     list of str
         A list with all available masks names.
+
     """
     return sorted(_available_masks.keys())
 
 
-def get_mask(
+def get_mask(  # noqa: C901
     masks: Union[str, Dict, List[Union[Dict, str]]],
     target_data: Dict[str, Any],
     extra_input: Optional[Dict[str, Any]] = None,
@@ -173,16 +288,49 @@ def get_mask(
     -------
     Nifti1Image
         The mask image.
+
+    Raises
+    ------
+    RuntimeError
+        If warp / transformation file extension is not ".mat" or ".h5" or
+        if fetch_icbm152_brain_gm_mask is used and requires warping to
+        other template space.
+    ValueError
+        If extra key is provided in addition to mask name in ``masks`` or
+        if no mask is provided or
+        if ``masks = "inherit"`` but ``extra_input`` is None or ``mask_item``
+        is None or ``mask_items``'s value is not in ``extra_input`` or
+        if callable parameters are passed to non-callable mask or
+        if parameters are passed to :func:`nilearn.masking.intersect_masks`
+        when there is only one mask or
+        if ``extra_input`` is None when ``target_data``'s space is native.
+
     """
+    # Check pre-requirements for space manipulation
+    target_space = target_data["space"]
+    # Set target standard space to target space
+    target_std_space = target_space
+    # Extra data type requirement check if target space is native
+    if target_space == "native":
+        # Check for extra inputs
+        if extra_input is None:
+            raise_error(
+                "No extra input provided, requires `Warp` and `T1w` "
+                "data types in particular for transformation to "
+                f"{target_data['space']} space for further computation."
+            )
+        # Set target standard space to warp file space source
+        target_std_space = extra_input["Warp"]["src"]
+
     # Get the min of the voxels sizes and use it as the resolution
     target_img = target_data["data"]
-    inherited_mask_item = target_data.get("mask_item", None)
     resolution = np.min(target_img.header.get_zooms()[:3])
 
+    # Convert masks to list if not already
     if not isinstance(masks, list):
         masks = [masks]
 
-    # Check that dicts have only one key
+    # Check that masks passed as dicts have only one key
     invalid_elements = [
         x for x in masks if isinstance(x, dict) and len(x) != 1
     ]
@@ -209,56 +357,157 @@ def get_mask(
 
     if len(true_masks) == 0:
         raise_error("No mask was passed. At least one mask is required.")
+
+    # Get the data type for the input data type's mask
+    inherited_mask_item = target_data.get("mask_item", None)
+
+    # Create component-scoped tempdir
+    tempdir = WorkDirManager().get_tempdir(prefix="masks")
+    # Create element-scoped tempdir so that warped mask is
+    # available later as nibabel stores file path reference for
+    # loading on computation
+    element_tempdir = WorkDirManager().get_element_tempdir(prefix="masks")
+
     # Get all the masks
     all_masks = []
     for t_mask in true_masks:
         if isinstance(t_mask, dict):
-            mask_name = list(t_mask.keys())[0]
+            mask_name = next(iter(t_mask.keys()))
             mask_params = t_mask[mask_name]
         else:
             mask_name = t_mask
             mask_params = None
 
+        # If mask is being inherited from previous steps like preprocessing
         if mask_name == "inherit":
+            # Requires extra input to be passed
             if extra_input is None:
                 raise_error(
                     "Cannot inherit mask from another data item "
                     "because no extra data was passed."
                 )
+            # Missing inherited mask item
             if inherited_mask_item is None:
                 raise_error(
                     "Cannot inherit mask from another data item "
                     "because no mask item was specified "
                     "(missing `mask_item` key in the data object)."
                 )
+            # Missing inherited mask item in extra input
             if inherited_mask_item not in extra_input:
                 raise_error(
                     "Cannot inherit mask from another data item "
                     f"because the item ({inherited_mask_item}) does not exist."
                 )
             mask_img = extra_input[inherited_mask_item]["data"]
+        # Starting with new mask
         else:
-            mask_object, _ = load_mask(
+            # Restrict fetch_icbm152_brain_gm_mask if target std space doesn't
+            # match
+            if (
+                mask_name == "fetch_icbm152_brain_gm_mask"
+                and target_std_space != "MNI152NLin2009aAsym"
+            ):
+                raise_error(
+                    (
+                        "``fetch_icbm152_brain_gm_mask`` is deprecated and "
+                        "space transformation to any other template space is "
+                        "prohibited as it will lead to unforeseen errors. "
+                        "``compute_brain_mask`` is a better alternative."
+                    ),
+                    klass=RuntimeError,
+                )
+            # Load mask
+            mask_object, _, mask_space = load_mask(
                 mask_name, path_only=False, resolution=resolution
             )
+            # Replace mask space with target space if mask's space is inherit
+            if mask_space == "inherit":
+                mask_space = target_std_space
+            # If mask is callable like from nilearn
             if callable(mask_object):
                 if mask_params is None:
                     mask_params = {}
-                mask_img = mask_object(target_img, **mask_params)
-            else:  # Mask is a Nifti1Image
+                # From nilearn
+                if mask_name != "compute_brain_mask":
+                    mask_img = mask_object(target_img, **mask_params)
+                # Not from nilearn
+                else:
+                    mask_img = mask_object(target_data, **mask_params)
+            # Mask is a Nifti1Image
+            else:
+                # Mask params provided
                 if mask_params is not None:
+                    # Unused params
                     raise_error(
                         "Cannot pass callable params to a non-callable mask."
                     )
+                # Resample mask to target image
                 mask_img = resample_to_img(
-                    mask_object,
-                    target_img,
+                    source_img=mask_object,
+                    target_img=target_img,
                     interpolation="nearest",
                     copy=True,
                 )
+            # Convert mask space if required
+            if mask_space != target_std_space:
+                # Get xfm file
+                xfm_file_path = get_xfm(src=mask_space, dst=target_std_space)
+                # Get target standard space template
+                target_std_space_template_img = get_template(
+                    space=target_std_space,
+                    target_data=target_data,
+                    extra_input=extra_input,
+                )
+
+                # Save mask image to a component-scoped tempfile
+                mask_path = tempdir / f"{mask_name}.nii.gz"
+                nib.save(mask_img, mask_path)
+
+                # Save template
+                target_std_space_template_path = (
+                    tempdir / f"{target_std_space}_T1w_{resolution}.nii.gz"
+                )
+                nib.save(
+                    target_std_space_template_img,
+                    target_std_space_template_path,
+                )
+
+                # Set warped mask path
+                warped_mask_path = element_tempdir / (
+                    f"{mask_name}_warped_from_{mask_space}_to_"
+                    f"{target_std_space}.nii.gz"
+                )
+
+                logger.debug(
+                    f"Using ANTs to warp {mask_name} "
+                    f"from {mask_space} to {target_std_space}"
+                )
+                # Set antsApplyTransforms command
+                apply_transforms_cmd = [
+                    "antsApplyTransforms",
+                    "-d 3",
+                    "-e 3",
+                    "-n 'GenericLabel[NearestNeighbor]'",
+                    f"-i {mask_path.resolve()}",
+                    f"-r {target_std_space_template_path.resolve()}",
+                    f"-t {xfm_file_path.resolve()}",
+                    f"-o {warped_mask_path.resolve()}",
+                ]
+                # Call antsApplyTransforms
+                run_ext_cmd(
+                    name="antsApplyTransforms", cmd=apply_transforms_cmd
+                )
+
+                mask_img = nib.load(warped_mask_path)
+
         all_masks.append(mask_img)
+
+    # Multiple masks, need intersection / union
     if len(all_masks) > 1:
+        # Intersect / union of masks
         mask_img = intersect_masks(all_masks, **intersect_params)
+    # Single mask
     else:
         if len(intersect_params) > 0:
             # Yes, I'm this strict!
@@ -268,20 +517,79 @@ def get_mask(
             )
         mask_img = all_masks[0]
 
-    return mask_img
+    # Warp mask if target data is native
+    if target_space == "native":
+        # Save mask image to a component-scoped tempfile
+        prewarp_mask_path = tempdir / "prewarp_mask.nii.gz"
+        nib.save(mask_img, prewarp_mask_path)
+
+        # Create an element-scoped tempfile for warped output
+        warped_mask_path = element_tempdir / "mask_warped.nii.gz"
+
+        # Check for warp file type to use correct tool
+        warp_file_ext = extra_input["Warp"]["path"].suffix
+        if warp_file_ext == ".mat":
+            logger.debug("Using FSL for mask warping")
+            # Set applywarp command
+            applywarp_cmd = [
+                "applywarp",
+                "--interp=nn",
+                f"-i {prewarp_mask_path.resolve()}",
+                # use resampled reference
+                f"-r {target_data['reference_path'].resolve()}",
+                f"-w {extra_input['Warp']['path'].resolve()}",
+                f"-o {warped_mask_path.resolve()}",
+            ]
+            # Call applywarp
+            run_ext_cmd(name="applywarp", cmd=applywarp_cmd)
+
+        elif warp_file_ext == ".h5":
+            logger.debug("Using ANTs for mask warping")
+            # Set antsApplyTransforms command
+            apply_transforms_cmd = [
+                "antsApplyTransforms",
+                "-d 3",
+                "-e 3",
+                "-n 'GenericLabel[NearestNeighbor]'",
+                f"-i {prewarp_mask_path.resolve()}",
+                # use resampled reference
+                f"-r {target_data['reference_path'].resolve()}",
+                f"-t {extra_input['Warp']['path'].resolve()}",
+                f"-o {warped_mask_path.resolve()}",
+            ]
+            # Call antsApplyTransforms
+            run_ext_cmd(name="antsApplyTransforms", cmd=apply_transforms_cmd)
+
+        else:
+            raise_error(
+                msg=(
+                    "Unknown warp / transformation file extension: "
+                    f"{warp_file_ext}"
+                ),
+                klass=RuntimeError,
+            )
+
+        # Load nifti
+        mask_img = nib.load(warped_mask_path)
+
+    # Delete tempdir
+    WorkDirManager().delete_tempdir(tempdir)
+
+    return mask_img  # type: ignore
 
 
 def load_mask(
     name: str,
     resolution: Optional[float] = None,
     path_only: bool = False,
-) -> Tuple[Optional[Union["Nifti1Image", Callable]], Optional[Path]]:
-    """Load mask.
+) -> Tuple[Optional[Union["Nifti1Image", Callable]], Optional[Path], str]:
+    """Load a mask.
 
     Parameters
     ----------
     name : str
-        The name of the mask.
+        The name of the mask. Check valid options by calling
+        :func:`.list_masks`.
     resolution : float, optional
         The desired resolution of the mask to load. If it is not
         available, the closest resolution will be loaded. Preferably, use a
@@ -296,16 +604,27 @@ def load_mask(
         Loaded mask image.
     pathlib.Path or None
         File path to the mask image.
+    str
+        The space of the mask.
+
+    Raises
+    ------
+    ValueError
+        If the ``name`` is invalid of if the mask family is invalid.
+
     """
-    mask_img = None
+    # Check for valid mask name
     if name not in _available_masks:
         raise_error(
-            f"Mask {name} not found. " f"Valid options are: {list_masks()}"
+            f"Mask {name} not found. Valid options are: {list_masks()}"
         )
 
+    # Copy mask definition to avoid edits in original object
     mask_definition = _available_masks[name].copy()
     t_family = mask_definition.pop("family")
 
+    # Check if the mask family is custom or built-in
+    mask_img = None
     if t_family == "CustomUserMask":
         mask_fname = Path(mask_definition["path"])
     elif t_family == "Vickery-Patil":
@@ -316,12 +635,16 @@ def load_mask(
     else:
         raise_error(f"I don't know about the {t_family} mask family.")
 
+    # Load mask
     if mask_fname is not None:
-        logger.info(f"Loading mask {mask_fname.absolute()}")
+        logger.info(f"Loading mask {mask_fname.absolute()!s}")
         if path_only is False:
+            # Load via nibabel
             mask_img = nib.load(mask_fname)
 
-    return mask_img, mask_fname
+    # Type-cast to remove error
+    mask_img = typing.cast("Nifti1Image", mask_img)
+    return mask_img, mask_fname, mask_definition["space"]
 
 
 def _load_vickery_patil_mask(
@@ -332,7 +655,7 @@ def _load_vickery_patil_mask(
 
     Parameters
     ----------
-    name : str
+    name : {"GM_prob0.2", "GM_prob0.2_cortex"}
         The name of the mask.
     resolution : float, optional
         The desired resolution of the mask to load. If it is not
@@ -344,6 +667,13 @@ def _load_vickery_patil_mask(
     -------
     pathlib.Path
         File path to the mask image.
+
+    Raises
+    ------
+    ValueError
+        If ``name`` is invalid or if ``resolution`` is invalid for
+        ``name = "GM_prob0.2"``.
+
     """
     if name == "GM_prob0.2":
         available_resolutions = [1.5, 3.0]
@@ -356,12 +686,14 @@ def _load_vickery_patil_mask(
             mask_fname = "CAT12_IXI555_MNI152_TMP_GS_GMprob0.2_clean.nii.gz"
         else:
             raise_error(
-                f"Cannot find a GM_prob0.2 mask for resolution {resolution}"
+                f"Cannot find a GM_prob0.2 mask of resolution {resolution}"
             )
     elif name == "GM_prob0.2_cortex":
         mask_fname = "GMprob0.2_cortex_3mm_NA_rm.nii.gz"
     else:
         raise_error(f"Cannot find a Vickery-Patil mask called {name}")
+
+    # Set path for masks
     mask_fname = _masks_path / "vickery-patil" / mask_fname
 
     return mask_fname