diff --git a/docs/changes/newsfragments/404.change b/docs/changes/newsfragments/404.change new file mode 100644 index 000000000..f4cf0187b --- /dev/null +++ b/docs/changes/newsfragments/404.change @@ -0,0 +1 @@ +Add ``source`` parameter and resurrect ``extra_input`` parameter for ``compute_brain_mask`` by `Fede Raimondo`_ diff --git a/docs/changes/newsfragments/404.enh b/docs/changes/newsfragments/404.enh new file mode 100644 index 000000000..726cdcfe2 --- /dev/null +++ b/docs/changes/newsfragments/404.enh @@ -0,0 +1 @@ +Enable ``compute_brain_mask`` to use subject-specific probseg files by `Fede Raimondo`_ diff --git a/junifer/data/masks/_masks.py b/junifer/data/masks/_masks.py index a603f75e3..b37df104b 100644 --- a/junifer/data/masks/_masks.py +++ b/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: