docs/changes/newsfragments/404.change

@@ -0,0 +1 @@
+Add ``source`` parameter and resurrect ``extra_input`` parameter for ``compute_brain_mask`` by `Fede Raimondo`_

docs/changes/newsfragments/404.enh

@@ -0,0 +1 @@
+Enable ``compute_brain_mask`` to use subject-specific probseg files by `Fede Raimondo`_

junifer/data/masks/_masks.py

@@ -47,6 +47,8 @@ def compute_brain_mask(
     warp_data: Optional[dict[str, Any]] = None,
     mask_type: str = "brain",
     threshold: float = 0.5,
+    source: str = "template",
+    extra_input: Optional[dict[str, Any]] = None,
 ) -> "Nifti1Image":
     """Compute the whole-brain, grey-matter or white-matter mask.
 
@@ -72,6 +74,13 @@ def compute_brain_mask(
         (default "brain").
     threshold : float, optional
         The value under which the template is cut off (default 0.5).
+    source : {"subject", "template"}, optional
+        The source of the mask. If "subject", the mask is computed from the
+        subject's data (``VBM_GM`` or ``VBM_WM``). If "template", the mask is
+        computed from the template data (default "template").
+    extra_input : dict, optional
+         The other fields in the data object. Useful for accessing other data
+         types (default None).
 
     Returns
     -------
@@ -82,7 +91,13 @@ def compute_brain_mask(
     ------
     ValueError
         If ``mask_type`` is invalid or
-        if ``warp_data`` is None when ``target_data``'s space is native.
+        if ``source`` is invalid or
+        if ``source="subject"`` and ``mask_type`` is invalid or
+        if ``warp_data`` is None when ``target_data``'s space is native or
+        if ``extra_input`` is None when ``source="subject"`` or
+        if ``VBM_GM`` or ``VBM_WM`` data types are not in ``extra_input``
+        when ``source="subject"`` and ``mask_type`` is ``"gm"`` or ``"wm"``
+        respectively.
 
     """
     logger.debug(f"Computing {mask_type} mask")
@@ -90,6 +105,12 @@ def compute_brain_mask(
     if mask_type not in ["brain", "gm", "wm"]:
         raise_error(f"Unknown mask type: {mask_type}")
 
+    if source not in ["subject", "template"]:
+        raise_error(f"Unknown mask source: {source}")
+
+    if source == "subject" and mask_type not in ["gm", "wm"]:
+        raise_error(f"Unknown mask type: {mask_type} for subject space")
+
     # Check pre-requirements for space manipulation
     if target_data["space"] == "native":
         # Warp data check
@@ -101,25 +122,50 @@ def compute_brain_mask(
         # Set space to fetch template using
         target_std_space = target_data["space"]
 
-    # Fetch template in closest resolution
-    template = get_template(
-        space=target_std_space,
-        target_img=target_data["data"],
-        extra_input=None,
-        template_type=mask_type,
-    )
-
-    # Resample and warp template if target space is native
-    if target_data["space"] == "native":
-        resampled_template = ANTsMaskWarper().warp(
-            mask_name=f"template_{target_std_space}_for_compute_brain_mask",
-            # use template here
-            mask_img=template,
-            src=target_std_space,
-            dst="native",
-            target_data=target_data,
-            warp_data=warp_data,
+    if source == "subject":
+        key = f"VBM_{mask_type.upper()}"
+        # Check for extra inputs
+        if extra_input is None:
+            raise_error(
+                f"No extra input provided, requires `{key}` "
+                "data type to infer target template data and space."
+            )
+        # Check for missing data type
+        if key not in extra_input:
+            raise_error(
+                f"Cannot compute {mask_type} from subject's data. "
+                f"Missing {key} in extra input."
+            )
+        template = extra_input[key]["data"]
+        template_space = extra_input[key]["space"]
+    else:
+        # Fetch template in closest resolution
+        template = get_template(
+            space=target_std_space,
+            target_img=target_data["data"],
+            extra_input=None,
+            template_type=mask_type,
         )
+        template_space = target_std_space
+    # Resample and warp template if target space is native
+    if target_data["space"] == "native" and template_space != "native":
+        if warp_data["warper"] == "fsl":
+            resampled_template = FSLMaskWarper().warp(
+                mask_name=f"template_{target_std_space}_for_compute_brain_mask",
+                mask_img=template,
+                target_data=target_data,
+                warp_data=warp_data,
+            )
+        elif warp_data["warper"] == "ants":
+            resampled_template = ANTsMaskWarper().warp(
+                mask_name=f"template_{target_std_space}_for_compute_brain_mask",
+                # use template here
+                mask_img=template,
+                src=target_std_space,
+                dst="native",
+                target_data=target_data,
+                warp_data=warp_data,
+            )
     # Resample template to target image
     else:
         resampled_template = resample_to_img(
@@ -503,7 +549,10 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
                     # custom compute_brain_mask
                     elif mask_name == "compute_brain_mask":
                         mask_img = mask_object(
-                            target_data, warper_spec, **mask_params
+                            target_data=target_data,
+                            warp_data=warper_spec,
+                            extra_input=extra_input,
+                            **mask_params,
                         )
                     # custom registered; arm kept for clarity
                     else: