#!/usr/bin/env python3
"""Cli interface."""
import json
import logging
import re
from enum import StrEnum
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Any
if TYPE_CHECKING:
from nilearn.maskers import NiftiMasker
import numpy as np
import typer
from numpy.typing import NDArray
from rich.logging import RichHandler
from rich.progress import Progress, SpinnerColumn, TextColumn
from patch_denoise import __version__
from patch_denoise.bindings.utils import (
DENOISER_MAP,
DENOISER_NAMES,
fast_cuda_check,
load_as_array,
load_complex_nifti,
save_array,
)
GPU_AVAILABLE = fast_cuda_check()
log = logging.getLogger(__name__)
logging.basicConfig(
level=logging.WARNING,
format="%(message)s",
handlers=[
RichHandler(show_time=False, show_path=True, show_level=True, markup=True)
],
)
logging.captureWarnings(True)
DENOISER_NAMES_HELP = ", ".join(DENOISER_NAMES)
[docs]
class AnalysisEnum(StrEnum):
"""Enum for BIDS analysis levels."""
PARTICIPANT = "participant"
[docs]
class DenoiserEnum(StrEnum):
"""Enum for denoising methods."""
MP_PCA = "mp-pca"
HYBRID_PCA = "hybrid-pca"
RAW = "raw"
OPTIMAL_FRO = "optimal-fro"
OPTIMAL_FRO_NOISE = "optimal-fro-noise"
OPTIMAL_NUC = "optimal-nuc"
OPTIMAL_OPE = "optimal-ope"
NORDIC = "nordic"
ADAPTIVE_QUT = "adaptive-qut"
[docs]
class RecombinationEnum(StrEnum):
"""Enum for recombination methods."""
WEIGHTED = "weighted"
MEAN = AVERAGE = "mean"
[docs]
def parse_dims(value: Any) -> tuple[int, ...]:
"""Parse a string representing dimensions into a 3 or 4-tuple of integers."""
if not isinstance(value, str):
return value # Already a tuple of ints
dims = [int(x) for x in re.findall(r"-?\d+", value)]
if len(dims) in (1, 3, 4):
return tuple(dims)
raise typer.BadParameter(
"Must be an int, 3-tuple, or 4-tuple ('11' or '11x11x11')"
" any 1-character separator is allowed (except space and -): "
"('11x11x11', 11_11_11', '11,11,11')"
)
[docs]
def parse_mask_arg(value: str):
"""Validate if value is 'auto' or points to an existing file."""
if value == "auto":
return value
path = Path(value)
if not path.exists() or not path.is_file():
raise typer.BadParameter(
f"Path should point to a file, or be 'auto': <{value}>."
)
return path.absolute()
app = typer.Typer(help="Patch denoising CLI tool.")
###########################
## Shared Argument Types ##
###########################
MethodOpt = Annotated[
DenoiserEnum,
typer.Option(
"-m", "--method", help=f"Denoising Method: Available: {DENOISER_NAMES_HELP}"
),
]
PatchShapeOpt = Annotated[
str,
typer.Option(
"-ps",
"--patch-shape",
parser=parse_dims,
metavar="X,Y,Z[,T]",
help="Patch shape. If 4D a sliding window is used. "
"If -1 is specified for a dimension, the entire dimension is put the patch.",
),
]
PatchOverlapOpt = Annotated[
str,
typer.Option(
"-po",
"--patch-overlap",
parser=parse_dims,
metavar="X,Y,Z[,T]",
help="Patch overlap. If 4D a sliding window is used. "
"If -1 is specified for a dimension, the entire dimension is put the patch.",
),
]
RecombinationOpt = Annotated[
RecombinationEnum,
typer.Option("-r", "--recombination", help="Recombination method."),
]
MaskOpt = Annotated[
str,
typer.Option(
"-k",
"--mask",
callback=parse_mask_arg,
help="Mask NIfTI file (3D). if auto, mask is computed automatically.",
metavar="MASK_FILE | auto",
),
]
MaskThreshOpt = Annotated[
int,
typer.Option(
"-t",
"--mask-threshold",
help="Min % of overlap between a patch and the mask to trigger computation.",
),
]
ExtraOpts = Annotated[
list[str] | None,
typer.Option(
"-e",
"--extra",
help="Extra parameters for the denoising method, passed as key=value pairs. "
"For example: --extra param1=val1 --extra param2=val2",
metavar="KEY=VALUE",
),
]
NaN2NumOpt = Annotated[
float | None, typer.Option(help="Replace any NaN in input-data with VALUE")
]
VerboseOpt = Annotated[
int,
typer.Option(
"-v", "--verbose", count=True, help="Increase verbosity level (e.g., -vvv)."
),
]
GpuFlag = Annotated[
bool,
typer.Option(
"--gpu/--cpu",
help="Use GPU or CPU for computation. Requires patch_denoise.gpu module. "
"GPU is enabled by default if available.",
),
]
#############
# Main CLI #
#############
def _load_noise_std(
noise_std_map_file: Path | None,
noise_std_map_phase_file: Path | None,
) -> tuple[NDArray | None, NDArray | None]:
if noise_std_map_file is not None and noise_std_map_phase_file is not None:
noise_std_map, affine_noise_map = load_complex_nifti(
noise_std_map_file,
noise_std_map_phase_file,
)
elif noise_std_map_file is not None:
noise_std_map, affine_noise_map = load_as_array(noise_std_map_file)
elif noise_std_map_phase_file is not None:
raise ValueError(
"The phase component of the noise map has been provided, "
"but not the magnitude."
)
else:
noise_std_map = None
affine_noise_map = None
return noise_std_map, affine_noise_map
def _load_validate_input(
input_file: Path,
input_phase: Path | None,
mask: Path | str,
noise_std_map_file: Path | None,
noise_std_map_phase_file: Path | None,
nan_to_num: float | None,
verbose: int,
) -> tuple[NDArray, NDArray, "NiftiMasker", NDArray | None]:
from nilearn.image import resample_img
from nilearn.maskers import NiftiMasker # noqa: F811 (runtime import)
if input_phase is not None:
input_data, affine = load_complex_nifti(input_file, input_phase)
else:
input_data, affine = load_as_array(input_file)
if nan_to_num is not None:
input_data = np.nan_to_num(input_data, nan=nan_to_num)
log.info(f"Input data shape: {input_data.shape}")
n_nans = np.isnan(input_data).sum()
if n_nans > 0:
log.warning(
f"{n_nans}/{input_data.size} voxels are NaN. "
"You might want to use --nan-to-num=<value>",
stacklevel=0,
)
masker = NiftiMasker(verbose=verbose, mask_strategy="epi")
if mask != "auto":
masker.mask_img = mask
masker.fit(input_file)
affine_mask = masker.mask_img_.affine
noise_std_map, affine_noise = _load_noise_std(
noise_std_map_file, noise_std_map_phase_file
)
if affine is not None:
if (affine_mask is not None) and not np.allclose(affine, affine_mask):
log.warning(
"Affine matrix of input and mask does not match, it will be resampled",
stacklevel=2,
)
masker.mask_img_ = resample_img(
masker.mask_img_,
target_affine=affine,
target_shape=input_data.shape[:3],
interpolation="nearest",
)
if (affine_noise is not None) and not np.allclose(affine, affine_noise):
log.warning(
"Affine matrix of input and noise map does not match", stacklevel=2
)
return input_data, affine, masker, noise_std_map
[docs]
@app.command()
def main(
input_file: Annotated[
Path,
typer.Argument(
exists=True,
dir_okay=False,
resolve_path=True,
help="Input noisy NIfTI file (4D).",
),
],
output_file: Annotated[
Path | None,
typer.Argument(
dir_okay=False,
resolve_path=True,
help="Output denoised NIfTI file (4D). Default is D<input_file>.",
),
] = None,
output_noise_std_map_file: Annotated[
Path | None,
typer.Option(
"--output-noise-std-map",
dir_okay=False,
resolve_path=True,
help="Output noise level estimation NIfTI file (3D).",
),
] = None,
method: MethodOpt = DenoiserEnum.OPTIMAL_FRO,
patch_shape: PatchShapeOpt = "11,11,11,-1",
patch_overlap: PatchOverlapOpt = "5,5,5,-1",
recombination: RecombinationOpt = RecombinationEnum.WEIGHTED,
mask: MaskOpt = "auto",
mask_threshold: MaskThreshOpt = 50,
extras: ExtraOpts = None,
nan_to_num: NaN2NumOpt = None,
verbose: VerboseOpt = 0,
gpu: GpuFlag = GPU_AVAILABLE,
input_phase: Annotated[
Path | None,
typer.Option(
"-ip",
"--input-phase",
exists=True,
dir_okay=False,
resolve_path=True,
help="Input phase NIfTI file (4D). If provided, process complex data.",
),
] = None,
noise_std_map_file: Annotated[
Path | None,
typer.Option(
"--noise-std-map",
exists=True,
dir_okay=False,
resolve_path=True,
help="Input Noise std map",
),
] = None,
noise_std_map_phase_file: Annotated[
Path | None,
typer.Option(
"--noise-std-map-phase",
exists=True,
dir_okay=False,
resolve_path=True,
help="Input Noise std map, phase component.",
),
] = None,
):
"""Perform local-low-rank denoising on 4D MRI data."""
kwargs = parse_extra_args(extras)
levels = [logging.WARNING, logging.INFO, logging.DEBUG]
level = levels[min(verbose, len(levels) - 1)]
logging.getLogger("patch_denoise").setLevel(level)
logging.getLogger("py.warnings").setLevel(level)
if output_file is None:
output_file = input_file.parent / f"D{input_file.name}"
parent_dir = output_file.parent
if not output_file.parent.exists():
parent_dir.mkdir(exist_ok=True, parents=True)
log.info(f"{output_file.parent} created")
if output_file.exists():
log.warning(f"{output_file} will be overwritten")
if output_noise_std_map_file is not None:
parent_dir = output_noise_std_map_file.parent
if not output_noise_std_map_file.parent.exists():
parent_dir.mkdir(exist_ok=True, parents=True)
log.info(f"{output_noise_std_map_file.parent} created")
if output_noise_std_map_file.exists():
log.warning(f"{output_noise_std_map_file} will be overwritten")
# 1. Define only the columns you want (just the spinner and the text)
with Progress(
SpinnerColumn(spinner_name="dots"),
TextColumn("[progress.description]{task.description}"),
) as progress:
progress.add_task(description="Loading and validating input...", total=None)
# 2. Add your task and start the progress display
input_data, affine, masker, noise_std_map = _load_validate_input(
input_file,
input_phase,
mask,
noise_std_map_file,
noise_std_map_phase_file,
nan_to_num,
verbose,
)
if mask == "auto":
mask_filename = output_file.with_stem("mask_" + output_file.stem)
log.info("Saving automatically computed mask to {mask_filename}.")
masker.mask_img_.to_filename(mask_filename)
log.info("Creating report for NiftiMasker.")
report = masker.generate_report()
report.save_as_html(output_file.with_suffix(".html"))
mask_data = masker.mask_img_.get_fdata().astype(bool)
from patch_denoise.space_time.base import _patch_param
# substitute any -1 in patch_shape or patch_overlap with the corresponding dimension
# of input_data
patch_shape_ = _patch_param(patch_shape, input_data.shape)
patch_overlap_ = _patch_param(patch_overlap, input_data.shape)
log.info(f"denoising method: {method}.")
log.info(f"patch shape: {patch_shape_} (from {patch_shape}).")
log.info(f"patch overlap: {patch_overlap_} (from {patch_overlap}).")
log.info(f"recombination method: {recombination}.")
log.info(f"mask threshold: {mask_threshold}.")
log.info(f"GPU: {gpu}.")
log.info(f"extra parameters: {kwargs}.")
log.info(f"nan_to_num: {nan_to_num}.")
log.info(f"input data shape: {input_data.shape}.")
log.info(msg=f"mask shape: {masker.mask_img_.shape}.")
log.info(
f"noise std map: {noise_std_map.shape if noise_std_map is not None else None}."
)
log.info(f"output file: {output_file}.")
log.info(f"output noise std map file: {output_noise_std_map_file}.")
log.debug(f"input affine:\n{affine}.")
log.debug(f"mask affine: \n{masker.mask_img_.affine}.")
if gpu:
if method not in [
DenoiserEnum.MP_PCA,
DenoiserEnum.OPTIMAL_FRO,
DenoiserEnum.OPTIMAL_FRO_NOISE,
DenoiserEnum.OPTIMAL_NUC,
DenoiserEnum.OPTIMAL_OPE,
]:
raise ValueError(f"Method {method} is not supported on GPU. ")
if not GPU_AVAILABLE:
raise RuntimeError(
"GPU support is not available. Please ensure that the "
"patch_denoise.gpu module is installed and that you have"
"a compatible GPU."
)
log.info("Using GPU for computation.")
from patch_denoise.gpu.main import main_gpu as denoise_func
kwargs["method"] = method
else:
denoise_func = DENOISER_MAP[method]
if method in [
DenoiserEnum.NORDIC,
DenoiserEnum.HYBRID_PCA,
DenoiserEnum.ADAPTIVE_QUT,
DenoiserEnum.OPTIMAL_FRO_NOISE,
]:
if noise_std_map is None:
raise RuntimeError("A noise map must be specified for this method.")
kwargs["noise_std"] = noise_std_map
denoised_data, _, noise_std_map, _ = denoise_func(
input_data,
patch_shape=patch_shape_,
patch_overlap=patch_overlap_,
mask=mask_data,
mask_threshold=mask_threshold,
recombination=recombination,
**kwargs,
)
save_array(denoised_data, affine, output_file)
if output_noise_std_map_file is not None:
save_array(noise_std_map, affine, output_noise_std_map_file)
############
# BIDS CLI #
############
bidsapp = typer.Typer(help="Patch denoising-bids CLI.")
[docs]
@bidsapp.command()
def bids_main(
bids_dir: Annotated[
Path,
typer.Argument(
exists=True,
file_okay=False,
dir_okay=True,
resolve_path=True,
help="Input BIDS directory.",
),
],
output_dir: Annotated[
Path,
typer.Argument(
exists=False,
file_okay=False,
dir_okay=True,
resolve_path=True,
help="Output BIDS directory.",
),
],
analysis_level: Annotated[
AnalysisEnum,
typer.Argument(
help="BIDS analysis level. Only 'participant' is supported.",
),
] = AnalysisEnum.PARTICIPANT,
participant_label: Annotated[
list[int] | None,
typer.Option(
"-participant-label",
"--participant_label",
help="List of participant labels to process. "
"If not provided, all participants will be processed.",
),
] = None,
session_label: Annotated[
list[int] | None,
typer.Option(
"-session-label",
"--session_label",
help="List of session labels to process."
" If not provided, all sessions will be processed.",
),
] = None,
task_label: Annotated[
list[str] | None,
typer.Option(
"-task-label",
"--task_label",
help="List of task labels to process."
" If not provided, all tasks will be processed.",
),
] = None,
bids_filters: Annotated[
Path | None,
typer.Option(
"-bids-filters",
"--bids_filters",
help="Path to a JSON file containing BIDS filters. See "
"https://fmriprep.org/en/latest/faq.html#"
"how-do-i-select-only-certain-files-to-be-input-to-fmriprep",
),
] = None,
method: MethodOpt = DenoiserEnum.OPTIMAL_FRO,
patch_shape: PatchShapeOpt = "11,11,11,-1",
patch_overlap: PatchOverlapOpt = "5,5,5,-1",
recombination: RecombinationOpt = RecombinationEnum.WEIGHTED,
mask: MaskOpt = "auto",
mask_threshold: MaskThreshOpt = 50,
extras: ExtraOpts = None,
nan_to_num: NaN2NumOpt = None,
verbose: VerboseOpt = 0,
gpu: GpuFlag = GPU_AVAILABLE,
noise_std_map_file: Annotated[
Path | None,
typer.Option(
"--noise-std-map",
exists=True,
dir_okay=False,
resolve_path=True,
help="Input Noise std map",
),
] = None,
noise_std_map_phase_file: Annotated[
Path | None,
typer.Option(
"--noise-std-map-phase",
exists=True,
dir_okay=False,
resolve_path=True,
help="Input Noise std map, phase component.",
),
] = None,
):
"""Run CLI for bids app."""
kwargs = parse_extra_args(extras)
levels = [logging.WARNING, logging.INFO, logging.DEBUG]
level = levels[min(verbose, len(levels) - 1)] # cap to last level index
log.setLevel(level=level)
if participant_label:
all_subjects = participant_label
else:
all_subjects = [
x.name.strip("sub-") for x in Path(bids_dir).iterdir() if "sub-" in x.name
]
filters = []
if task_label:
filters.append(("task", task_label[0]))
if session_label:
filters.append(("ses", session_label[0]))
noise_std_map, _ = _load_noise_std(noise_std_map_file, noise_std_map_phase_file)
if gpu and not GPU_AVAILABLE:
raise RuntimeError(
"GPU support is not available. Please ensure that the "
"patch_denoise.gpu module is installed and that you have"
"a compatible GPU."
)
if method in [
DenoiserEnum.NORDIC,
DenoiserEnum.HYBRID_PCA,
DenoiserEnum.ADAPTIVE_QUT,
DenoiserEnum.OPTIMAL_FRO_NOISE,
]:
if noise_std_map is None:
raise RuntimeError("A noise map must be specified for this method.")
kwargs["noise_std"] = noise_std_map
if gpu:
from patch_denoise.gpu.main import main_gpu as denoise_func
kwargs["method"] = method
else:
denoise_func = DENOISER_MAP[method]
from nilearn.image import load_img
from nilearn.interfaces.bids import get_bids_files, parse_bids_filename
from nilearn.interfaces.bids.utils import bids_entities, create_bids_filename
from nilearn.maskers import NiftiMasker
from patch_denoise.space_time.base import _patch_param
output_dir.mkdir(exist_ok=True, parents=True)
ds_json = output_dir / "dataset_description.json"
GeneratedBy = {
"Name": "patch-denoise",
"Version": __version__,
"Description": ("A dataset denoised with patch-denoise."),
"CodeURL": "https://github.com/paquiteau/patch-denoising",
}
if ds_json.exists():
with ds_json.open() as f_obj:
dataset_description = json.load(f_obj)
if dataset_description.get("GeneratedBy"):
dataset_description["GeneratedBy"].append(GeneratedBy)
else:
dataset_description = {
"BIDSVersion": "1.9.0",
"DatasetType": "derivative",
"GeneratedBy": [GeneratedBy],
}
with ds_json.open("w") as f_obj:
json.dump(dataset_description, f_obj, indent=4, sort_keys=True)
entities_to_include = bids_entities()["raw"] + bids_entities()["derivatives"]
for sub_label in all_subjects:
files = get_bids_files(
bids_dir,
file_tag="bold",
modality_folder="func",
file_type="nii.*",
sub_label=sub_label,
filters=filters,
)
for f in files:
parsed_f = parse_bids_filename(f)
output_dir_for_subject = output_dir / f"sub-{parsed_f['entities']['sub']}"
if "ses" in parsed_f["entities"]:
output_dir_for_subject = (
output_dir_for_subject / f"ses-{parsed_f['entities']['ses']}"
)
parsed_f["entities"]["desc"] = "denoised"
output_filename = output_dir_for_subject / create_bids_filename(
parsed_f,
entities_to_include=entities_to_include,
)
parsed_f["suffix"] = "mask"
output_mask_filename = output_dir_for_subject / create_bids_filename(
parsed_f,
entities_to_include=entities_to_include,
)
parsed_f["suffix"] = "std"
output_std_filename = output_dir_for_subject / create_bids_filename(
parsed_f,
entities_to_include=entities_to_include,
)
output_dir_for_subject.mkdir(exist_ok=True, parents=True)
masker = NiftiMasker(verbose=verbose, mask_strategy="epi")
masker.fit(f)
mask = masker.mask_img_.get_fdata().astype(bool)
affine = masker.mask_img_.affine
report = masker.generate_report()
report.save_as_html(output_mask_filename.with_suffix(".html"))
masker.mask_img_.to_filename(output_mask_filename)
input_data = load_img(f).get_fdata()
patch_shape_ = _patch_param(patch_shape, input_data.shape)
patch_overlap_ = _patch_param(patch_overlap, input_data.shape)
denoised_data, _, noise_std_map, _ = denoise_func(
input_data,
patch_shape=patch_shape_,
patch_overlap=patch_overlap_,
mask=mask,
mask_threshold=mask_threshold,
recombination=recombination,
**kwargs,
)
print(noise_std_map.shape)
save_array(denoised_data, affine, output_filename)
save_array(noise_std_map, affine, output_std_filename)
if __name__ == "__main__":
app()