junifer 0.0.4.dev694__py3-none-any.whl → 0.0.4.dev781__py3-none-any.whl

junifer/data/masks.py

@@ -20,16 +20,16 @@ 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 ..pipeline import WorkDirManager
-from ..utils import logger, raise_error, run_ext_cmd
+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
 
 
@@ -40,10 +40,91 @@ if TYPE_CHECKING:
 _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
@@ -59,7 +140,20 @@ 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
@@ -123,7 +217,7 @@ def register_mask(
     mask_path : str or pathlib.Path
         The path to the mask file.
     space : str
-        The space of the mask.
+        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).
@@ -198,30 +292,45 @@ def get_mask(  # noqa: C901
     Raises
     ------
     RuntimeError
-        If masks are in different spaces and they need to be intersected /
-        unionized or
-        if warp / transformation file extension is not ".mat" or ".h5".
+        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 multiple masks are provided and their spaces do not match 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
     ]
@@ -248,9 +357,19 @@ def get_mask(  # noqa: C901
 
     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 = []
-    all_spaces = []
     for t_mask in true_masks:
         if isinstance(t_mask, dict):
             mask_name = next(iter(t_mask.keys()))
@@ -281,21 +400,40 @@ def get_mask(  # noqa: C901
                     f"because the item ({inherited_mask_item}) does not exist."
                 )
             mask_img = extra_input[inherited_mask_item]["data"]
-            mask_space = target_data["space"]
         # Starting with new mask
         else:
+            # 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_data["space"]
+                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)
+                # 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
@@ -306,31 +444,69 @@ def get_mask(  # noqa: C901
                     )
                 # 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,
                 )
-        all_spaces.append(mask_space)
+            # 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:
-        # Make a set of unique spaces
-        unique_spaces = set(all_spaces)
-        # Intersect / union of masks only if all masks are in the same space
-        if len(unique_spaces) == 1:
-            mask_img = intersect_masks(all_masks, **intersect_params)
-            # Store the mask space for further checks
-            mask_space = next(iter(unique_spaces))
-        else:
-            raise_error(
-                msg=(
-                    f"Masks are in different spaces: {unique_spaces}, "
-                    "unable to merge."
-                ),
-                klass=RuntimeError,
-            )
+        # Intersect / union of masks
+        mask_img = intersect_masks(all_masks, **intersect_params)
     # Single mask
     else:
         if len(intersect_params) > 0:
@@ -340,30 +516,13 @@ def get_mask(  # noqa: C901
                 "when there is only one mask."
             )
         mask_img = all_masks[0]
-        mask_space = all_spaces[0]
-
-    # Warp mask if target data is native and mask space is not native
-    if target_data["space"] == "native" and target_data["space"] != mask_space:
-        # 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."
-            )
-
-        # Create component-scoped tempdir
-        tempdir = WorkDirManager().get_tempdir(prefix="masks")
 
+    # 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 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")
-
         # Create an element-scoped tempfile for warped output
         warped_mask_path = element_tempdir / "mask_warped.nii.gz"
 
@@ -413,8 +572,8 @@ def get_mask(  # noqa: C901
         # Load nifti
         mask_img = nib.load(warped_mask_path)
 
-        # Delete tempdir
-        WorkDirManager().delete_tempdir(tempdir)
+    # Delete tempdir
+    WorkDirManager().delete_tempdir(tempdir)
 
     return mask_img  # type: ignore
 

junifer/data/parcellations.py

@@ -22,6 +22,7 @@ from nilearn import datasets, image
 
 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
 
 
@@ -154,7 +155,7 @@ def register_parcellation(
     parcels_labels : list of str
         The list of labels for the parcellation.
     space : str
-        The space of the parcellation.
+        The template space of the parcellation, for e.g., "MNI152NLin6Asym".
     overwrite : bool, optional
         If True, overwrite an existing parcellation with the same name.
         Does not apply to built-in parcellations (default False).
@@ -236,78 +237,126 @@ def get_parcellation(
     Raises
     ------
     RuntimeError
-        If parcellations are in different spaces and they need to be merged or
-        if warp / transformation file extension is not ".mat" or ".h5".
+        If warp / transformation file extension is not ".mat" or ".h5".
     ValueError
         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"]
     resolution = np.min(target_img.header.get_zooms()[:3])
 
+    # Create component-scoped tempdir
+    tempdir = WorkDirManager().get_tempdir(prefix="parcellations")
+    # Create element-scoped tempdir so that warped parcellation is
+    # available later as nibabel stores file path reference for
+    # loading on computation
+    element_tempdir = WorkDirManager().get_element_tempdir(
+        prefix="parcellations"
+    )
+
     # Load the parcellations
     all_parcellations = []
     all_labels = []
-    all_spaces = []
     for name in parcellation:
         img, labels, _, space = load_parcellation(
             name=name,
             resolution=resolution,
         )
-        # Resample all of them to the image
-        resampled_img = image.resample_to_img(
+
+        # Convert parcellation spaces if required
+        if space != target_std_space:
+            # Get xfm file
+            xfm_file_path = get_xfm(src=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 parcellation image to a component-scoped tempfile
+            parcellation_path = tempdir / f"{name}.nii.gz"
+            nib.save(img, parcellation_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 parcellation path
+            warped_parcellation_path = element_tempdir / (
+                f"{name}_warped_from_{space}_to_" f"{target_std_space}.nii.gz"
+            )
+
+            logger.debug(
+                f"Using ANTs to warp {name} "
+                f"from {space} to {target_std_space}"
+            )
+            # Set antsApplyTransforms command
+            apply_transforms_cmd = [
+                "antsApplyTransforms",
+                "-d 3",
+                "-e 3",
+                "-n 'GenericLabel[NearestNeighbor]'",
+                f"-i {parcellation_path.resolve()}",
+                f"-r {target_std_space_template_path.resolve()}",
+                f"-t {xfm_file_path.resolve()}",
+                f"-o {warped_parcellation_path.resolve()}",
+            ]
+            # Call antsApplyTransforms
+            run_ext_cmd(name="antsApplyTransforms", cmd=apply_transforms_cmd)
+
+            img = nib.load(warped_parcellation_path)
+
+        # Resample parcellation to target image
+        img_to_merge = image.resample_to_img(
             source_img=img,
             target_img=target_img,
             interpolation="nearest",
             copy=True,
         )
-        all_parcellations.append(resampled_img)
+
+        all_parcellations.append(img_to_merge)
         all_labels.append(labels)
-        all_spaces.append(space)
 
     # Avoid merging if there is only one parcellation
     if len(all_parcellations) == 1:
         resampled_parcellation_img = all_parcellations[0]
         labels = all_labels[0]
+    # Parcellations are already transformed to target standard space
     else:
-        # Merge the parcellations only if all parcellations are in the same
-        # space
-        if len(set(all_spaces)) == 1:
-            resampled_parcellation_img, labels = merge_parcellations(
-                parcellations_list=all_parcellations,
-                parcellations_names=parcellation,
-                labels_lists=all_labels,
-            )
-        else:
-            raise_error(
-                msg="Parcellations are in different spaces, unable to merge.",
-                klass=RuntimeError,
-            )
-
-    # Warp parcellation if target data is native
-    if target_data["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."
-            )
-
-        # Create component-scoped tempdir
-        tempdir = WorkDirManager().get_tempdir(prefix="parcellations")
+        resampled_parcellation_img, labels = merge_parcellations(
+            parcellations_list=all_parcellations,
+            parcellations_names=parcellation,
+            labels_lists=all_labels,
+        )
 
+    # Warp parcellation if target space is native
+    if target_space == "native":
         # Save parcellation image to a component-scoped tempfile
         prewarp_parcellation_path = tempdir / "prewarp_parcellation.nii.gz"
         nib.save(resampled_parcellation_img, prewarp_parcellation_path)
 
-        # Create element-scoped tempdir so that warped parcellation is
-        # available later as nibabel stores file path reference for
-        # loading on computation
-        element_tempdir = WorkDirManager().get_element_tempdir(
-            prefix="parcellations"
-        )
         # Create an element-scoped tempfile for warped output
         warped_parcellation_path = (
             element_tempdir / "parcellation_warped.nii.gz"
@@ -359,8 +408,8 @@ def get_parcellation(
         # Load nifti
         resampled_parcellation_img = nib.load(warped_parcellation_path)
 
-        # Delete tempdir
-        WorkDirManager().delete_tempdir(tempdir)
+    # Delete tempdir
+    WorkDirManager().delete_tempdir(tempdir)
 
     return resampled_parcellation_img, labels  # type: ignore
 

junifer/data/template_spaces.py

@@ -99,6 +99,7 @@ def get_template(
     space: str,
     target_data: Dict[str, Any],
     extra_input: Optional[Dict[str, Any]] = None,
+    template_type: str = "T1w",
 ) -> nib.Nifti1Image:
     """Get template for the space, tailored for the target image.
 
@@ -112,6 +113,8 @@ def get_template(
     extra_input : dict, optional
         The other fields in the data object. Useful for accessing other data
         types (default None).
+    template_type : {"T1w", "brain", "gm", "wm", "csf"}, optional
+        The template type to retrieve (default "T1w").
 
     Returns
     -------
@@ -121,15 +124,19 @@ def get_template(
     Raises
     ------
     ValueError
-        If ``space`` is invalid.
+        If ``space`` or ``template_type`` is invalid.
     RuntimeError
-        If template in the required resolution is not found.
+        If required template is not found.
 
     """
     # Check for invalid space; early check to raise proper error
     if space not in tflow.templates():
         raise_error(f"Unknown template space: {space}")
 
+    # Check for template type
+    if template_type not in ["T1w", "brain", "gm", "wm", "csf"]:
+        raise_error(f"Unknown template type: {template_type}")
+
     # Get the min of the voxels sizes and use it as the resolution
     target_img = target_data["data"]
     resolution = np.min(target_img.header.get_zooms()[:3]).astype(int)
@@ -145,18 +152,38 @@ def get_template(
     logger.info(f"Downloading template {space} in resolution {resolution}")
     # Retrieve template
     try:
+        suffix = None
+        desc = None
+        label = None
+        if template_type == "T1w":
+            suffix = template_type
+            desc = None
+            label = None
+        elif template_type == "brain":
+            suffix = "mask"
+            desc = "brain"
+            label = None
+        elif template_type in ["gm", "wm", "csf"]:
+            suffix = "probseg"
+            desc = None
+            label = template_type.upper()
+        # Set kwargs for fetching
+        kwargs = {
+            "suffix": suffix,
+            "desc": desc,
+            "label": label,
+        }
         template_path = tflow.get(
             space,
             raise_empty=True,
             resolution=resolution,
-            suffix="T1w",
-            desc=None,
             extension="nii.gz",
+            **kwargs,
         )
     except Exception:  # noqa: BLE001
         raise_error(
-            f"Template {space} not found in the required resolution "
-            f"{resolution}",
+            f"Template {space} ({template_type}) with resolution {resolution} "
+            "not found",
             klass=RuntimeError,
         )
     else: