[ENH]: Compute GM/WM/CSF mask from probseg files #400

New issue

Closed

opened 2024-11-25 12:41:59 +00:00 by fraimondo · 0 comments

fraimondo commented

2024-11-25 12:41:59 +00:00

(Migrated from github.com)

Are you requiring a new dataset or marker?

I understand this is not a marker or dataset request

Which feature do you want to include?

The compute_brain_mask function allows to take a probseg file and apply a threshold.

Currently, this masks are from template spaces. However, some datasets (mostly when providing fmriprep derivatives) can have subject specific probsegs which should be preferred.

How do you imagine this integrated in junifer?

add a src parameter to the compute_brain_mask function which can be either template or subject. If template, get the template from templateflow, if `subject, look for the corresponding VBM (or MASK) in the subject's data.

Do you have a sample code that implements this outside of junifer?

def compute_brain_mask(
    target_data: dict[str, Any],
    warp_data: Optional[dict[str, Any]] = None,
    mask_type: str = "brain",
    threshold: float = 0.5,
    src: str = "template",
    extra_input: Optional[dict[str, Any]] = None,
) -> "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``. If target space is native, then the template is
    warped to native and then thresholded.

    Parameters
    ----------
    target_data : dict
        The corresponding item of the data object for which mask will be
        loaded.
    warp_data : dict or None, optional
        The warp data item of the data object. Needs to be provided if
        ``target_data`` is in native space (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 ``warp_data`` 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}")

    if src not in ["subject", "template"]:
        raise_error(f"Unknown mask source: {src}")

    if src == "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
        if warp_data is None:
            raise_error("No `warp_data` provided")
        # Set space to fetch template using
        target_std_space = warp_data["src"]
    else:
        # Set space to fetch template using
        target_std_space = target_data["space"]

    if src == "subject":
        key = f"VBM_{mask_type.upper()}"
        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_data=target_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(
            source_img=template, target_img=target_data["data"]
        )

    # Threshold resampled template and get mask
    mask = (get_data(resampled_template) >= threshold).astype("int8")

    return new_img_like(target_data["data"], mask)  # type: ignore

Anything else to say?

No response

### Are you requiring a new dataset or marker? - [X] I understand this is not a marker or dataset request ### Which feature do you want to include? The `compute_brain_mask` function allows to take a probseg file and apply a threshold. Currently, this masks are from template spaces. However, some datasets (mostly when providing fmriprep derivatives) can have subject specific probsegs which should be preferred. ### How do you imagine this integrated in junifer? add a `src` parameter to the `compute_brain_mask` function which can be either `template` or `subject`. If template, get the template from templateflow, if `subject, look for the corresponding VBM (or MASK) in the subject's data. ### Do you have a sample code that implements this outside of junifer? ```shell def compute_brain_mask( target_data: dict[str, Any], warp_data: Optional[dict[str, Any]] = None, mask_type: str = "brain", threshold: float = 0.5, src: str = "template", extra_input: Optional[dict[str, Any]] = None, ) -> "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``. If target space is native, then the template is warped to native and then thresholded. Parameters ---------- target_data : dict The corresponding item of the data object for which mask will be loaded. warp_data : dict or None, optional The warp data item of the data object. Needs to be provided if ``target_data`` is in native space (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 ``warp_data`` 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}") if src not in ["subject", "template"]: raise_error(f"Unknown mask source: {src}") if src == "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 if warp_data is None: raise_error("No `warp_data` provided") # Set space to fetch template using target_std_space = warp_data["src"] else: # Set space to fetch template using target_std_space = target_data["space"] if src == "subject": key = f"VBM_{mask_type.upper()}" 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_data=target_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( source_img=template, target_img=target_data["data"] ) # Threshold resampled template and get mask mask = (get_data(resampled_template) >= threshold).astype("int8") return new_img_like(target_data["data"], mask) # type: ignore ``` ### Anything else to say? _No response_