From acbfeef2fab6a8084f50752fac5da4398be28b62 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Fri, 29 May 2026 18:03:31 +0000 Subject: [PATCH 01/22] working on prostate --- examples/fastmri_inference_plot.py | 312 +++++--- examples/run_all.py | 684 +++++++++++++++++ mri_recon/distortions/utils.py | 140 ++++ .../examples/run_experiments_fastmri_brain.py | 313 ++++++++ .../run_experiments_fastmri_prostateT2.py | 289 ++++++++ mri_recon/reconstruction/inference.py | 17 + mri_recon/utils/__init__.py | 7 + mri_recon/utils/grappa.py | 215 ++++++ mri_recon/utils/oasis_adapter.py | 163 +++- mri_recon/utils/plot.py | 21 +- mri_recon/utils/prostate_adaptor.py | 461 ++++++++++++ mri_recon/utils/recon_prostate_T2.py | 700 ++++++++++++++++++ pyproject.toml | 1 + uv.lock | 145 ++-- 14 files changed, 3277 insertions(+), 191 deletions(-) create mode 100644 examples/run_all.py create mode 100644 mri_recon/distortions/utils.py create mode 100644 mri_recon/examples/run_experiments_fastmri_brain.py create mode 100644 mri_recon/examples/run_experiments_fastmri_prostateT2.py create mode 100644 mri_recon/utils/grappa.py create mode 100644 mri_recon/utils/prostate_adaptor.py create mode 100644 mri_recon/utils/recon_prostate_T2.py diff --git a/examples/fastmri_inference_plot.py b/examples/fastmri_inference_plot.py index 1379db0..b7a1913 100644 --- a/examples/fastmri_inference_plot.py +++ b/examples/fastmri_inference_plot.py @@ -35,11 +35,11 @@ ) from mri_recon.reconstruction import ( ConjugateGradientReconstructor, - EXPLICIT_UNET_ALGORITHMS, OASISSinglecoilUnetReconstructor, choose_reconstructor, uses_oasis_centered_path, validate_algorithm_dataset_compatibility, + EXPLICIT_UNET_ALGORITHMS, ) from mri_recon.utils import ( OasisCenteredFFTPhysics, @@ -52,49 +52,54 @@ ) FASTMRI_REPORT_DIR = Path("reports") / "fastmri_inference_plot" +FASTMRI_MULTICOIL_REPORT_DIR = Path("reports") / "fastmri_multicoil_inference_plot" OASIS_REPORT_DIR = Path("reports") / "oasis_inference_plot" +CMRXRECON_REPORT_DIR = Path("reports") / "cmrxrecon_inference_plot" FASTMRI_REPORT_DIR.mkdir(parents=True, exist_ok=True) +FASTMRI_MULTICOIL_REPORT_DIR.mkdir(parents=True, exist_ok=True) OASIS_REPORT_DIR.mkdir(parents=True, exist_ok=True) +CMRXRECON_REPORT_DIR.mkdir(parents=True, exist_ok=True) ALGORITHMS = [ - "zero-filled", - # "conjugate-gradient", + # "zero-filled", + "conjugate-gradient", # "ram", # "dip", - "tv-pgd", + # "tv-pgd", # "wavelet-fista", - "tv-fista", + # "tv-fista", # "tv-pdhg", *list(EXPLICIT_UNET_ALGORITHMS), ] DISTORTIONS = [ - "Cartesian undersampling (variable density)", - "Cartesian undersampling (uniform random)", - "Cartesian undersampling (uniform random, zero ACS)", - "Cartesian undersampling (equispaced)", + "no distortion", + # "Cartesian undersampling (variable density)", + # "Cartesian undersampling (uniform random)", + # "Cartesian undersampling (uniform random, zero ACS)", + # "Cartesian undersampling (equispaced)", "Cartesian undersampling (equispaced, zero ACS)", - "Partial Fourier", - "Phase-encode ghosting", - "Segmented translation motion", - "Segmented rotational motion", - "Translation motion", - "Rotational motion", - "Off-center anisotropic Gaussian bias field", - "Gaussian bias field", - "Anisotropic LP", - "Hann taper LP", - "Kaiser taper LP", - "Gaussian noise", - "Isotropic LP", - "Radial high-pass emphasis", + # "Partial Fourier", + # "Phase-encode ghosting", + # "Segmented translation motion", + # "Segmented rotational motion", + # "Translation motion", + # "Rotational motion", + # "Off-center anisotropic Gaussian bias field", + # "Gaussian bias field", + # "Anisotropic LP", + # "Hann taper LP", + # "Kaiser taper LP", + # "Gaussian noise", + # "Isotropic LP", + # "Radial high-pass emphasis", ] METRICS = [ "PSNR", - "NMSE", - "SSIM", - "HaarPSI", - "SharpnessIndex", - "BlurStrength", + # "NMSE", + # "SSIM", + # "HaarPSI", + # "SharpnessIndex", + # "BlurStrength", ] @@ -211,6 +216,8 @@ def choose_distortion( return GaussianKspaceBiasField(width_fraction=0.35, edge_gain=0.4) case "Gaussian noise": return GaussianNoiseDistortion(sigma=0.00001) + case "no distortion": + return BaseDistortion() case _: raise ValueError(f"Unknown distortion {name!r}") @@ -248,16 +255,48 @@ def prepare_measurement_sample( """ if dataset_name == "oasis": - reference_image = sample_batch["x"].to(run_device) - return reference_image, image_to_kspace(reference_image) - - # FastMRI batches are tuples such as (x, y) or (x, y, params). - y_fastmri = sample_batch[1].to(run_device) - if use_oasis_fft_path: - reference_image = fastmri_measurement_to_image(y_fastmri, device=run_device) - return reference_image, fastmri_measurement_to_oasis_kspace(y_fastmri, device=run_device) + x = sample_batch["x"].to(run_device) + y = image_to_kspace(x) + coil_maps = None + elif dataset_name in ("fastmri",) and use_oasis_fft_path: + y = sample_batch[1].to(run_device) + x = fastmri_measurement_to_image(y) + y = fastmri_measurement_to_oasis_kspace(y, device=run_device) + coil_maps = None + elif dataset_name == "fastmri_multicoil" and use_oasis_fft_path: + y = sample_batch[1].to(run_device) + + coil_maps = ( + sample_batch[2]["coil_maps"].to(run_device) + if isinstance(sample_batch, (tuple, list)) + and len(sample_batch) == 3 + and "coil_maps" in sample_batch[2] + else None + ) + x = fastmri_measurement_to_image(y, coil_maps=coil_maps) + y = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) + elif dataset_name in ("fastmri", "fastmri_multicoil"): + x = None + y = sample_batch[1].to(run_device) + coil_maps = ( + sample_batch[2]["coil_maps"].to(run_device) + if isinstance(sample_batch, (tuple, list)) + and len(sample_batch) == 3 + and "coil_maps" in sample_batch[2] + else None + ) + elif dataset_name in ("cmrxrecon"): + x = sample_batch[0].to(run_device) + y = sample_batch[1].to(run_device) + coil_maps = ( + sample_batch[2]["coil_maps"].to(run_device) + if isinstance(sample_batch, (tuple, list)) + and len(sample_batch) == 3 + and "coil_maps" in sample_batch[2] + else None + ) - return None, y_fastmri + return x, y, coil_maps def build_physics_pair( @@ -265,6 +304,7 @@ def build_physics_pair( distortion_operator: BaseDistortion, run_device: torch.device | str, use_oasis_fft_path: bool, + coil_maps: torch.Tensor | None = None, ) -> tuple[object, object]: """Build clean and distorted physics operators for the active path.""" @@ -276,11 +316,13 @@ def build_physics_pair( clean_physics = DistortedKspaceMultiCoilMRI( distortion=BaseDistortion(), img_size=(1, 2, *image_shape), + coil_maps=coil_maps, device=run_device, ) distorted_physics = DistortedKspaceMultiCoilMRI( distortion=distortion_operator, img_size=(1, 2, *image_shape), + coil_maps=coil_maps, device=run_device, ) return clean_physics, distorted_physics @@ -295,7 +337,11 @@ def build_physics_pair( type=Path, help="Local FastMRI directory with raw k-space .h5 files or OASIS root directory.", ) - parser.add_argument("--dataset", choices=("fastmri", "oasis"), default="fastmri") + parser.add_argument( + "--dataset", + choices=("fastmri", "oasis", "fastmri_multicoil", "cmrxrecon"), + default="fastmri", + ) parser.add_argument("--distortion", type=str, default="", choices=DISTORTIONS) parser.add_argument( @@ -334,7 +380,16 @@ def build_physics_pair( validate_algorithm_dataset_compatibility(args.dataset, algo_name) # set up report dir - REPORT_DIR = OASIS_REPORT_DIR if args.dataset == "oasis" else FASTMRI_REPORT_DIR + if args.dataset == "fastmri": + REPORT_DIR = FASTMRI_REPORT_DIR + elif args.dataset == "oasis": + REPORT_DIR = OASIS_REPORT_DIR + elif args.dataset == "fastmri_multicoil": + REPORT_DIR = FASTMRI_MULTICOIL_REPORT_DIR + elif args.dataset == "cmrxrecon": + REPORT_DIR = CMRXRECON_REPORT_DIR + else: + raise NotImplementedError(f"Invalid dataset: {args.dataset}") # set up device, dataset, metrics device = dinv.utils.get_device() @@ -345,8 +400,23 @@ def build_physics_pair( split_csv=split_csv, sample_rate=0.6, ) - else: + elif args.dataset == "fastmri": dataset = dinv.datasets.FastMRISliceDataset(str(args.source), slice_index="middle") + elif args.dataset == "fastmri_multicoil": + dataset = dinv.datasets.FastMRISliceDataset( + str(args.source), + slice_index="middle", + transform=dinv.datasets.MRISliceTransform( + estimate_coil_maps=True, + acs=15, + ), + ) + elif args.dataset == "cmrxrecon": + dataset = dinv.datasets.CMRxReconSliceDataset( + str(args.source), data_dir="SingleCoil/Cine/TrainingSet/FullSample", apply_mask=False + ) + else: + raise NotImplementedError(f"Invalid dataset: {args.dataset}") metrics = [choose_metric(m) for m in METRICS] for i, batch in enumerate(iter(torch.utils.data.DataLoader(dataset))): @@ -355,83 +425,91 @@ def build_physics_pair( break for algo_name in selected_algorithms: - use_oasis_path = uses_oasis_centered_path(args.dataset, algo_name) - x_reference, y = prepare_measurement_sample( - sample_batch=batch, - dataset_name=args.dataset, - use_oasis_fft_path=use_oasis_path, - run_device=device, - ) - algo = choose_reconstructor( - algo_name, - img_size=y.shape[-2:], - device=device, - verbose=args.verbose, - dataset=args.dataset, - ).to(device) - - for distortion_name in selected_distortions: - distortion = choose_distortion( - distortion_name, - keep_fraction=args.keep_fraction, - center_fraction=args.center_fraction, - cartesian_axis=-1 if use_oasis_path else -2, - ) - - physics_clean, physics = build_physics_pair( - image_shape=y.shape[-2:], - distortion_operator=distortion, - run_device=device, + try: + use_oasis_path = uses_oasis_centered_path(args.dataset, algo_name) + x_reference, y, coil_maps = prepare_measurement_sample( + sample_batch=batch, + dataset_name=args.dataset, use_oasis_fft_path=use_oasis_path, + run_device=device, ) - y_distorted = distortion.A(y) - - # generate reference reconstructions (CG) for both clean and distorted k-space - # without correction for the distortion, i.e. using physics_clean in both cases - if use_oasis_path: - x_clean = x_reference - x_distorted = kspace_to_image(y_distorted) - else: - x_clean = ConjugateGradientReconstructor()(y, physics_clean) - x_distorted = ConjugateGradientReconstructor()(y_distorted, physics_clean) - - save_kspace_plot( - y, - y_distorted, - REPORT_DIR / f"DISTORTION_{algo_name}_{distortion_name}_sample_{i}.png", - distortion_name, - ) - - print(f"Evaluating algo {algo_name}, distortion {distortion_name}, sample {i}...") - - # actual reconstruction with the algo being evaluated - x_uncorrected = algo(y_distorted, physics_clean) - x_corrected = algo(y_distorted, physics) - - print("done!") - - dinv.utils.plot( - { - "Undistorted ksp, CG recon": x_clean, - "Distorted ksp, CG recon": x_distorted, - f"Distorted ksp, {algo_name} recon, uncorrected": x_uncorrected, - f"Distorted ksp, {algo_name} recon, corrected": x_corrected, - }, - subtitles=[ - "", - "", - "\n".join( - f"{m.__class__.__name__} {m(x_uncorrected, x_clean).item():.2f}" - for m in metrics - ), - "\n".join( - f"{m.__class__.__name__} {m(x_corrected, x_clean).item():.2f}" - for m in metrics - ), - ], - show=False, - close=True, - suptitle=f"Algo {algo_name}, distortion {distortion_name}, Sample {i}", - save_fn=REPORT_DIR / f"ALGO_{algo_name}_{distortion_name}_sample_{i}.png", - fontsize=3, + algo = choose_reconstructor( + algo_name, + img_size=y.shape[-2:], + device=device, + verbose=args.verbose, + dataset=args.dataset, + ).to(device) + + for distortion_name in selected_distortions: + distortion = choose_distortion( + distortion_name, + keep_fraction=args.keep_fraction, + center_fraction=args.center_fraction, + cartesian_axis=-1 if use_oasis_path else -2, + ) + + physics_clean, physics = build_physics_pair( + image_shape=y.shape[-2:], + distortion_operator=distortion, + run_device=device, + use_oasis_fft_path=use_oasis_path, + coil_maps=coil_maps, + ) + y_distorted = distortion.A(y) + + # generate reference reconstructions (CG) for both clean and distorted k-space + # without correction for the distortion, i.e. using physics_clean in both cases + if use_oasis_path: + x_clean = x_reference + x_distorted = kspace_to_image(y_distorted) + else: + x_clean = ConjugateGradientReconstructor()(y, physics_clean) + x_distorted = ConjugateGradientReconstructor()(y_distorted, physics_clean) + + save_kspace_plot( + y, + y_distorted, + REPORT_DIR / f"DISTORTION_{algo_name}_{distortion_name}_sample_{i}.png", + distortion_name, + ) + + print( + f"Evaluating algo {algo_name}, distortion {distortion_name}, sample {i}..." + ) + + # actual reconstruction with the algo being evaluated + x_uncorrected = algo(y_distorted, physics_clean) + x_corrected = algo(y_distorted, physics) + + print("done!") + + dinv.utils.plot( + { + "Undistorted ksp, CG recon": x_clean, + "Distorted ksp, CG recon": x_distorted, + f"Distorted ksp, {algo_name} recon, uncorrected": x_uncorrected, + f"Distorted ksp, {algo_name} recon, corrected": x_corrected, + }, + subtitles=[ + "", + "", + "\n".join( + f"{m.__class__.__name__} {m(x_uncorrected, x_clean).item():.2f}" + for m in metrics + ), + "\n".join( + f"{m.__class__.__name__} {m(x_corrected, x_clean).item():.2f}" + for m in metrics + ), + ], + show=False, + close=True, + suptitle=f"Algo {algo_name}, distortion {distortion_name}, Sample {i}", + save_fn=REPORT_DIR / f"ALGO_{algo_name}_{distortion_name}_sample_{i}.png", + fontsize=3, + ) + except Exception as e: + print( + f"Error processing algo {algo_name}, distortion {distortion_name}, sample {i}: {e}" ) diff --git a/examples/run_all.py b/examples/run_all.py new file mode 100644 index 0000000..205d171 --- /dev/null +++ b/examples/run_all.py @@ -0,0 +1,684 @@ +"""Inference various reconstructors for various distortion operators. + +Usage: + python examples/fastmri_inference_plot.py --source ../ram-experiments/data/fastmri/knee/singlecoil_val +""" + +import os +import sys + + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import numpy as np +from pathlib import Path +import deepinv as dinv +import torch +from tifffile import imwrite + +from mri_recon.distortions import ( + AnisotropicResolutionReduction, + BaseDistortion, + CartesianUndersampling, + DistortedKspaceMultiCoilMRI, + GaussianKspaceBiasField, + GaussianNoiseDistortion, + HannTaperResolutionReduction, + IsotropicResolutionReduction, + KaiserTaperResolutionReduction, + OffCenterAnisotropicGaussianKspaceBiasField, + PartialFourierDistortion, + PhaseEncodeGhostingDistortion, + RadialHighPassEmphasisDistortion, + RotationalMotionDistortion, + SegmentedRotationalMotionDistortion, + SegmentedTranslationMotionDistortion, + TranslationMotionDistortion, +) +from mri_recon.reconstruction import ( + ConjugateGradientReconstructor, + choose_reconstructor, + uses_oasis_centered_path, + validate_algorithm_dataset_compatibility, + EXPLICIT_UNET_ALGORITHMS, +) +from mri_recon.utils import ( + OasisCenteredFFTPhysics, + OasisCenterSliceFolderDataset, + FastMRIProstateDataset, + fastmri_measurement_to_image, + fastmri_measurement_to_oasis_kspace, + oasis_kspace_to_fastmri_measurement, + image_to_kspace, + _kspace_to_log_magnitude, +) + +EXPERIMENTS_DIR = Path("reports") / "experiments" + +ALGORITHMS = [ + # "zero-filled", + "conjugate-gradient", + # "ram", + # "dip", + # "tv-pgd", + # "wavelet-fista", + # "tv-fista", + # "tv-pdhg", + *list(EXPLICIT_UNET_ALGORITHMS), +] + +DISTORTIONS = [ + "no distortion", + # "Cartesian undersampling (variable density)", + # "Cartesian undersampling (uniform random)", + # "Cartesian undersampling (uniform random, zero ACS)", + # "Cartesian undersampling (equispaced)", + "Cartesian undersampling (equispaced, zero ACS)", + # "Partial Fourier", + # "Phase-encode ghosting", + # "Segmented translation motion", + # "Segmented rotational motion", + # "Translation motion", + # "Rotational motion", + # "Off-center anisotropic Gaussian bias field", + # "Gaussian bias field", + # "Anisotropic LP", + # "Hann taper LP", + # "Kaiser taper LP", + # "Gaussian noise", + # "Isotropic LP", + # "Radial high-pass emphasis", +] +METRICS = [ + "PSNR", + # "NMSE", + # "SSIM", + # "HaarPSI", + # "SharpnessIndex", + # "BlurStrength", +] + +DATASETS = { + # "fastmri": "/home/melanie.dohmen/mri_recon/data/fastmri/singlecoil_val", + "oasis": "/home/melanie.dohmen/mri_recon/data/oasis", + # "fastmri_multicoil": "/home/melanie.dohmen/mri_recon/data/fastmri/multicoil_train", + "cmrxrecon": "/home/melanie.dohmen/mri_recon/data/CMRxRecon/CMRxRecon/", # SingleCoil/Cine/TrainingSet/FullSample", + "prostate": "/home/melanie.dohmen/mri_recon/data/fastmri/fastMRI_prostate_T2_IDS_001_020", +} + + +def convert_image_for_save(im: torch.Tensor) -> np.ndarray: + """ + Convert a PyTorch tensor image complex tensor to a real-valued NumPy array suitable + by calculating the magnitude. + (B, 2, H, W) or (B, H, W) with complex type -> (B, H, W) + + Args: + im (torch.Tensor): The input image tensor. + + Returns: + np.ndarray: The converted image array. + """ + if torch.is_complex(im) or im.shape[1] == 2: + im = dinv.utils.signals.complex_abs(im, dim=1, keepdim=False) + return im.numpy() + + +def choose_distortion( + name: str, + keep_fraction: float = 0.25, + center_fraction: float = 0.125, + cartesian_axis: int = -2, +) -> BaseDistortion: + """Build one distortion operator for the inference comparison script. + + The ``cartesian_axis`` is supplied by the active measurement convention: + FastMRI-native runs use the repository's existing axis, while OASIS-native + and FastMRI-to-OASIS runs use the centered OASIS axis. + """ + + match name: + case "Phase-encode ghosting": + return PhaseEncodeGhostingDistortion( + line_period=2, + line_offset=1, + phase_error_radians=torch.pi / 2, + corrupted_line_scale=1.0, + ) + case "Cartesian undersampling (variable density)": + return CartesianUndersampling( + keep_fraction=keep_fraction, + center_fraction=center_fraction, + pattern="variable_density_random", + axis=cartesian_axis, + seed=42, + ) + case "Cartesian undersampling (uniform random)": + return CartesianUndersampling( + keep_fraction=keep_fraction, + center_fraction=center_fraction, + pattern="uniform_random", + axis=cartesian_axis, + seed=42, + ) + case "Cartesian undersampling (uniform random, zero ACS)": + return CartesianUndersampling( + keep_fraction=keep_fraction, + center_fraction=0.0, + pattern="uniform_random", + axis=cartesian_axis, + seed=42, + ) + case "Cartesian undersampling (equispaced)": + return CartesianUndersampling( + keep_fraction=keep_fraction, + center_fraction=center_fraction, + pattern="equispaced", + axis=cartesian_axis, + seed=42, + ) + case "Cartesian undersampling (equispaced, zero ACS)": + return CartesianUndersampling( + keep_fraction=keep_fraction, + center_fraction=0.0, + pattern="equispaced", + axis=cartesian_axis, + seed=42, + ) + case "Partial Fourier": + return PartialFourierDistortion( + partial_fraction=0.7, + center_fraction=center_fraction, + axis=cartesian_axis, + side="high", + ) + case "Anisotropic LP": + return AnisotropicResolutionReduction( + kx_radius_fraction=1.0, + ky_radius_fraction=0.25, + ) + case "Hann taper LP": + return HannTaperResolutionReduction( + radius_fraction=0.35, + transition_fraction=0.4, + ) + case "Kaiser taper LP": + return KaiserTaperResolutionReduction( + radius_fraction=0.35, + transition_fraction=0.4, + beta=8.6, + ) + case "Radial high-pass emphasis": + return RadialHighPassEmphasisDistortion(alpha=0.4) + case "Isotropic LP": + return IsotropicResolutionReduction(radius_fraction=0.1) + case "Off-center anisotropic Gaussian bias field": + return OffCenterAnisotropicGaussianKspaceBiasField( + width_x_fraction=0.2, + width_y_fraction=0.35, + center_x_fraction=0.15, + center_y_fraction=-0.1, + edge_gain=0.3, + ) + case "Translation motion": + return TranslationMotionDistortion(shift_x_pixels=60, shift_y_pixels=10) + case "Rotational motion": + return RotationalMotionDistortion(angle_radians=torch.pi / 6) + case "Segmented rotational motion": + return SegmentedRotationalMotionDistortion( + angle_radians=(0.0, torch.pi / 20, -torch.pi / 24, torch.pi / 16), + ) + case "Segmented translation motion": + return SegmentedTranslationMotionDistortion( + shift_x_pixels=(0.0, 20.0, 50.0, -50.0), + shift_y_pixels=(0.0, 10.0, -20.0, 20.0), + ) + case "Gaussian bias field": + return GaussianKspaceBiasField(width_fraction=0.35, edge_gain=0.4) + case "Gaussian noise": + return GaussianNoiseDistortion(sigma=0.00001) + case "no distortion": + return BaseDistortion() + case _: + raise ValueError(f"Unknown distortion {name!r}") + + +def choose_metric(name: str) -> dinv.metric.Metric: + """Build one evaluation metric used in the saved comparison plots.""" + + match name: + case "PSNR": + return dinv.metric.PSNR(max_pixel=None, complex_abs=True) + case "NMSE": + return dinv.metric.NMSE(complex_abs=True) + case "SSIM": + return dinv.metric.SSIM(max_pixel=None, complex_abs=True) + case "HaarPSI": + return dinv.metric.HaarPSI(norm_inputs="min_max", complex_abs=True) + case "BlurStrength": + return dinv.metric.BlurStrength(complex_abs=True) + case "SharpnessIndex": + return dinv.metric.SharpnessIndex(complex_abs=True) + + +# def prepare_measurement_sample( +# sample_batch: object, +# dataset_name: str, +# use_oasis_fft_path: bool, +# run_device: torch.device | str, +# ) -> tuple[torch.Tensor | None, torch.Tensor]: +# """Prepare one input measurement and its clean image reference. + +# FastMRI samples are loaded as native measurements. When the OASIS U-Net is +# selected on FastMRI data, the helper converts those measurements into the +# centered OASIS k-space convention while preserving the native adjoint image +# as the clean reference. +# """ + +# if dataset_name == "oasis": +# x = sample_batch["x"].to(run_device) +# y = image_to_kspace(x) +# coil_maps = None +# elif dataset_name in ("fastmri",) and use_oasis_fft_path: +# y = sample_batch[1].to(run_device) +# x = fastmri_measurement_to_image(y) +# y = fastmri_measurement_to_oasis_kspace(y, device=run_device) +# coil_maps = None +# elif dataset_name == "fastmri_multicoil" and use_oasis_fft_path: +# y = sample_batch[1].to(run_device) + +# coil_maps = ( +# sample_batch[2]["coil_maps"].to(run_device) +# if isinstance(sample_batch, (tuple, list)) +# and len(sample_batch) == 3 +# and "coil_maps" in sample_batch[2] +# else None +# ) +# x = fastmri_measurement_to_image(y, coil_maps=coil_maps) +# y = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) +# elif dataset_name in ("fastmri", "fastmri_multicoil"): +# x = None +# y = sample_batch[1].to(run_device) +# coil_maps = ( +# sample_batch[2]["coil_maps"].to(run_device) +# if isinstance(sample_batch, (tuple, list)) +# and len(sample_batch) == 3 +# and "coil_maps" in sample_batch[2] +# else None +# ) +# elif dataset_name in ("cmrxrecon"): +# x = sample_batch[0].to(run_device) +# y = sample_batch[1].to(run_device) +# coil_maps = ( +# sample_batch[2]["coil_maps"].to(run_device) +# if isinstance(sample_batch, (tuple, list)) +# and len(sample_batch) == 3 +# and "coil_maps" in sample_batch[2] +# else None +# ) +# elif dataset_name in ("prostate"): +# x = sample_batch[0].to(run_device) +# y = sample_batch[1].to(run_device) +# coil_maps = None + +# print(f"\t[Prepared measurement] k-space shape {y.shape} and reference image shape: {x.shape if x is not None else None}") + +# return x, y, coil_maps + + +def get_measurement_sample( + sample_batch: object, + dataset_name: str, + run_device: torch.device | str, +) -> tuple[torch.Tensor | None, torch.Tensor]: + """Prepare one input measurement and its clean image reference. + + FastMRI samples are loaded as native measurements. When the OASIS U-Net is + selected on FastMRI data, the helper converts those measurements into the + centered OASIS k-space convention while preserving the native adjoint image + as the clean reference. + """ + coil_maps = None + if dataset_name == "oasis": + x = sample_batch["x"].to(run_device) + print(f"\t[Debug] Reference image shape: {x.shape}, dtype: {x.dtype}") + y_centered = image_to_kspace(x) + print(f"\t[Debug] Centered k-space shape: {y_centered.shape}, dtype: {y_centered.dtype}") + y = oasis_kspace_to_fastmri_measurement(y_centered, device=run_device) + elif dataset_name in ("fastmri"): + y = sample_batch[1].to(run_device) + y_centered = fastmri_measurement_to_oasis_kspace(y, device=run_device) + x = fastmri_measurement_to_image(y, rss=True) + elif dataset_name in ("fastmri_multicoil"): + y = sample_batch[1].to(run_device) + coil_maps = ( + sample_batch[2]["coil_maps"].to(run_device) + if isinstance(sample_batch, (tuple, list)) + and len(sample_batch) == 3 + and "coil_maps" in sample_batch[2] + else None + ) + x = fastmri_measurement_to_image(y, coil_maps=coil_maps, rss=True) + y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) + elif dataset_name in ("cmrxrecon"): + # ignore multi-coil reference image data + x = sample_batch[0].to(run_device) + print(f"\t[Debug] Reference image shape: {x.shape}, dtype: {x.dtype}") + + y = sample_batch[1].to(run_device) + + coil_maps = ( + sample_batch[2]["coil_maps"].to(run_device) + if isinstance(sample_batch, (tuple, list)) + and len(sample_batch) == 3 + and "coil_maps" in sample_batch[2] + else None + ) + y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) + # reconstruct coil-combined image reference from multi-coil k-space data using + # integrated espirit sensitivity map estimation, RSS coil combination + x = fastmri_measurement_to_image(y, coil_maps=coil_maps, rss=True) + + elif dataset_name in ("prostate"): + # has shape: (B, num_averages, coils, H, W) with dtype= complex128-> take first average and convert to image space reference + x = sample_batch[0].to(run_device) + print(f"\t[Debug] Reference image shape: {x.shape}, type: {x.dtype}") + # take mean of average images: + x = x.mean(dim=1) + print(f"\t[Debug] Mean of averages image shape: {x.shape}, type: {x.dtype}") + + # convert to channel representation of complex numbers + # (B, H, W) with complex dtype -> (B, H, W, 2) with real dtype + x = torch.view_as_real(x) if torch.is_complex(x) else x + print(f"\t[Debug] after view_as_real (if complex) shape: {x.shape}, type: {x.dtype}") + + # move channel with real and imaginary parts to channel dimension + # (B, H, W, 2) -> (B, 2, H, W) + x = x.moveaxis(-1, 1) + print(f"\t[Debug] after moving channels: {x.shape}, type: {x.dtype}") + + # ignore k-space data: + y = sample_batch[1].to(run_device) + print(f"\t[Debug] Original k-space shape: {y.shape}, type: {y.dtype}") + # y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) + # create oasis-like k-space data from image: + y_centered = image_to_kspace(x) + print(f"\t[Debug] Centered k-space shape: {y_centered.shape}, type: {y_centered.dtype}") + y = oasis_kspace_to_fastmri_measurement(y_centered, device=run_device) + + print(f"\tk-space shape {y.shape} and reference image shape: {x.shape}") + + return x, y, y_centered, coil_maps + + +if __name__ == "__main__": + # parser = argparse.ArgumentParser(description=__doc__) + + # # data related arguments + # parser.add_argument( + # "--source", + # type=Path, + # help="Local FastMRI directory with raw k-space .h5 files or OASIS root directory.", + # ) + # parser.add_argument( + # "--dataset", + # choices=("fastmri", "oasis", "fastmri_multicoil", "cmrxrecon"), + # default="fastmri", + # ) + + # parser.add_argument("--distortion", type=str, default="", choices=DISTORTIONS) + # parser.add_argument( + # "--keep_fraction", + # type=float, + # default=0.25, + # help="Fraction of k-space lines to keep for undersampling distortions.", + # ) + # parser.add_argument( + # "--center_fraction", + # type=float, + # default=0.125, + # help="Fraction of low-frequency k-space lines to keep fully for undersampling distortions.", + # ) + + # # algo related arguments + # parser.add_argument( + # "--algorithm", + # type=str, + # default="", + # choices=ALGORITHMS, + # help="Reconstruction algorithm applied to undistorted and distorted k-space.", + # ) + # # inference related arguments + # parser.add_argument("--num_samples", type=int, default=1, help="How many samples to process.") + # parser.add_argument( + # "--verbose", + # action="store_true", + # help="Enable verbose output for reconstructors that support it.", + # ) + # args = parser.parse_args() + num_samples = 1 + keep_fraction = 0.25 + center_fraction = 0.125 + verbose = True + + os.makedirs(EXPERIMENTS_DIR, exist_ok=True) + + # set up device, dataset, metrics + device = dinv.utils.get_device() + + for dataset_name, dataset_rootdir in DATASETS.items(): + print(f"=== {dataset_name} ===") + + selected_algorithms = ALGORITHMS + selected_distortions = DISTORTIONS + + if dataset_name == "oasis": + # split_csv = OASISSinglecoilUnetReconstructor.resolve_default_split_csv() + dataset = OasisCenterSliceFolderDataset( + data_path=dataset_rootdir, + ) + elif dataset_name == "fastmri": + dataset = dinv.datasets.FastMRISliceDataset(str(dataset_rootdir), slice_index="middle") + elif dataset_name == "fastmri_multicoil": + dataset = dinv.datasets.FastMRISliceDataset( + str(dataset_rootdir), + slice_index="middle", + transform=dinv.datasets.MRISliceTransform( + estimate_coil_maps=True, + acs=15, + ), + ) + elif dataset_name == "cmrxrecon": + dataset = dinv.datasets.CMRxReconSliceDataset( + str(dataset_rootdir), + data_dir="SingleCoil/Cine/TrainingSet/FullSample", + apply_mask=False, + ) + elif dataset_name == "prostate": + dataset = FastMRIProstateDataset(data_path=dataset_rootdir, num_samples=num_samples) + else: + raise NotImplementedError(f"Invalid dataset: {dataset_name}") + metrics = [choose_metric(m) for m in METRICS] + + for i, batch in enumerate(iter(torch.utils.data.DataLoader(dataset))): + # exit loop if we have processed the specified number of samples + if i >= num_samples: + break + + print(f"{dataset_name} sample {i}...") + x_reference, y, y_centered, coil_maps = get_measurement_sample( + sample_batch=batch, + dataset_name=dataset_name, + run_device=device, + ) + + physics_clean_oasis_fft_path = OasisCenteredFFTPhysics(BaseDistortion()) + physics_clean_fastmri_path = DistortedKspaceMultiCoilMRI( + BaseDistortion(), img_size=x_reference.shape, coil_maps=coil_maps, device=device + ) + + x_clean_oasis_fft_path = ConjugateGradientReconstructor()( + y, physics_clean_oasis_fft_path + ) + x_clean_fastmri_path = ConjugateGradientReconstructor()(y, physics_clean_fastmri_path) + + # reference reconstructions: + imwrite( + os.path.join( + EXPERIMENTS_DIR, f"image_{dataset_name}_sample_{i}_CG_oasis_fft_path.tiff" + ), + convert_image_for_save(x_clean_oasis_fft_path), + ) + imwrite( + os.path.join( + EXPERIMENTS_DIR, f"image_{dataset_name}_sample_{i}_CG_fastmri_path.tiff" + ), + convert_image_for_save(x_clean_fastmri_path), + ) + imwrite( + os.path.join(EXPERIMENTS_DIR, f"image_{dataset_name}_sample_{i}_reference.tiff"), + convert_image_for_save(x_reference), + ) + + for distortion_name in selected_distortions: + print(f"\t{distortion_name} ...") + distortion_oasis_fft_path = choose_distortion( + distortion_name, + keep_fraction=keep_fraction, + center_fraction=center_fraction, + cartesian_axis=-1, + ) + + distortion_fastmri_path = choose_distortion( + distortion_name, + keep_fraction=keep_fraction, + center_fraction=center_fraction, + cartesian_axis=-2, + ) + + y_distorted_oasis_fft_path = distortion_oasis_fft_path.A(y_centered) + y_distorted_fastmri_path = distortion_fastmri_path.A(y) + + physics_distorted_oasis_fft_path = OasisCenteredFFTPhysics( + distortion_oasis_fft_path + ) + physics_distorted_fastmri_path = DistortedKspaceMultiCoilMRI( + distortion_fastmri_path, + img_size=x_reference.shape, + coil_maps=coil_maps, + device=device, + ) + + for algo_name in selected_algorithms: + print(f"\t\t{algo_name} ...") + try: + validate_algorithm_dataset_compatibility(dataset_name, algo_name) + + use_oasis_path = uses_oasis_centered_path(dataset_name, algo_name) + if use_oasis_path: + y_distorted = y_distorted_oasis_fft_path + physics_clean = physics_clean_oasis_fft_path + physics_distorted = physics_distorted_oasis_fft_path + x_clean = x_reference + else: + y_distorted = y_distorted_fastmri_path + physics_clean = physics_clean_fastmri_path + physics_distorted = physics_distorted_fastmri_path + + algo = choose_reconstructor( + algo_name, + img_size=y_distorted.shape[-2:], + device=device, + verbose=verbose, + dataset=dataset_name, + ).to(device) + + # save reference and distorted k-space for debugging and visualization purposes + imwrite( + os.path.join( + EXPERIMENTS_DIR, f"kspace_{dataset_name}_sample_{i}_reference.tiff" + ), + _kspace_to_log_magnitude(y).numpy(), + ) + imwrite( + os.path.join( + EXPERIMENTS_DIR, + f"kspace_{dataset_name}_sample_{i}_{distortion_name}.tiff", + ), + _kspace_to_log_magnitude(y_distorted).numpy(), + ) + + # actual reconstruction with the algo being evaluated + try: + if dataset_name == "prostate": + # prostate dataset has multiple k-space averages, + # so we reconstruct each average separately and then average in the image domain + x_corrected_averages = [] + x_uncorrected_averages = [] + for average in range(y_distorted.shape[0]): + x_uncorrected_averages.append( + algo(y_distorted[average], physics_clean) + ) + x_corrected_averages.append( + algo(y_distorted[average], physics_distorted) + ) + + x_uncorrected = torch.stack(x_uncorrected_averages, dim=0).mean( + dim=0 + ) + x_corrected = torch.stack(x_corrected_averages, dim=0).mean(dim=0) + + else: + x_uncorrected = algo(y_distorted, physics_clean) + x_corrected = algo(y_distorted, physics_distorted) + + # performed reconstruction images + imwrite( + os.path.join( + EXPERIMENTS_DIR, + f"image_{dataset_name}_sample_{i}_{distortion_name}_{algo_name}_uncorrected.tiff", + ), + convert_image_for_save(x_uncorrected), + ) + imwrite( + os.path.join( + EXPERIMENTS_DIR, + f"image_{dataset_name}_sample_{i}_{distortion_name}_{algo_name}_corrected.tiff", + ), + convert_image_for_save(x_corrected), + ) + + except Exception as e: + print( + f"Error reconstructing algo {algo_name} with distortion {distortion_name} on sample {i}: {e}" + ) + + # dinv.utils.plot( + # { + # "Undistorted ksp, CG recon": x_clean, + # "Distorted ksp, CG recon": x_distorted, + # f"Distorted ksp, {algo_name} recon, uncorrected": x_uncorrected, + # f"Distorted ksp, {algo_name} recon, corrected": x_corrected, + # }, + # subtitles=[ + # "", + # "", + # "\n".join( + # f"{m.__class__.__name__} {m(x_uncorrected, x_clean).item():.2f}" + # for m in metrics + # ), + # "\n".join( + # f"{m.__class__.__name__} {m(x_corrected, x_clean).item():.2f}" + # for m in metrics + # ), + # ], + # show=False, + # close=True, + # suptitle=f"Algo {algo_name}, distortion {distortion_name}, Sample {i}", + # save_fn=REPORT_DIR / f"ALGO_{algo_name}_{distortion_name}_sample_{i}.png", + # fontsize=3, + # ) + except Exception as e: + print( + f"\t\tError processing algo {algo_name}, distortion {distortion_name}, sample {i}: {e}" + ) diff --git a/mri_recon/distortions/utils.py b/mri_recon/distortions/utils.py new file mode 100644 index 0000000..96d00b7 --- /dev/null +++ b/mri_recon/distortions/utils.py @@ -0,0 +1,140 @@ +import torch + +from .base import BaseDistortion +from .resolution import ( + HannTaperResolutionReduction, + IsotropicResolutionReduction, + AnisotropicResolutionReduction, + KaiserTaperResolutionReduction, + RadialHighPassEmphasisDistortion, +) +from .undersampling import CartesianUndersampling, PartialFourierDistortion +from .biasfield import OffCenterAnisotropicGaussianKspaceBiasField, GaussianKspaceBiasField +from .noise import GaussianNoiseDistortion +from .motion import ( + RotationalMotionDistortion, + SegmentedRotationalMotionDistortion, + TranslationMotionDistortion, + SegmentedTranslationMotionDistortion, +) +from .ghosting import PhaseEncodeGhostingDistortion + + +def choose_distortion( + name: str, + keep_fraction: float = 0.25, + center_fraction: float = 0.125, + cartesian_axis: int = -2, +) -> BaseDistortion: + """Build one distortion operator for the inference comparison script. + + The ``cartesian_axis`` is supplied by the active measurement convention: + FastMRI-native runs use the repository's existing axis, while OASIS-native + and FastMRI-to-OASIS runs use the centered OASIS axis. + """ + + match name: + case "PhaseEncodeGhosting": + return PhaseEncodeGhostingDistortion( + line_period=2, + line_offset=1, + phase_error_radians=torch.pi / 2, + corrupted_line_scale=1.0, + ) + case "CartesianUndersamplingVariableDensity": + return CartesianUndersampling( + keep_fraction=keep_fraction, + center_fraction=center_fraction, + pattern="variable_density_random", + axis=cartesian_axis, + seed=42, + ) + case "CartesianUndersamplingUniformRandom": + return CartesianUndersampling( + keep_fraction=keep_fraction, + center_fraction=center_fraction, + pattern="uniform_random", + axis=cartesian_axis, + seed=42, + ) + case "CartesianUndersamplingUniformRandomZeroACS": + return CartesianUndersampling( + keep_fraction=keep_fraction, + center_fraction=0.0, + pattern="uniform_random", + axis=cartesian_axis, + seed=42, + ) + case "CartesianUndersamplingEquispaced": + return CartesianUndersampling( + keep_fraction=keep_fraction, + center_fraction=center_fraction, + pattern="equispaced", + axis=cartesian_axis, + seed=42, + ) + case "CartesianUndersamplingEquispacedZeroACS": + return CartesianUndersampling( + keep_fraction=keep_fraction, + center_fraction=0.0, + pattern="equispaced", + axis=cartesian_axis, + seed=42, + ) + case "PartialFourier": + return PartialFourierDistortion( + partial_fraction=0.7, + center_fraction=center_fraction, + axis=cartesian_axis, + side="high", + ) + case "AnisotropicLP": + return AnisotropicResolutionReduction( + kx_radius_fraction=1.0, + ky_radius_fraction=0.25, + ) + case "HannTaperLP": + return HannTaperResolutionReduction( + radius_fraction=0.35, + transition_fraction=0.4, + ) + case "KaiserTaperLP": + return KaiserTaperResolutionReduction( + radius_fraction=0.35, + transition_fraction=0.4, + beta=8.6, + ) + case "RadialHighPassEmphasis": + return RadialHighPassEmphasisDistortion(alpha=0.4) + case "IsotropicLP": + return IsotropicResolutionReduction(radius_fraction=0.1) + case "OffCenterAnisotropicGaussianKspaceBiasField": + return OffCenterAnisotropicGaussianKspaceBiasField( + width_x_fraction=0.2, + width_y_fraction=0.35, + center_x_fraction=0.15, + center_y_fraction=-0.1, + edge_gain=0.3, + ) + case "TranslationMotion": + return TranslationMotionDistortion(shift_x_pixels=60, shift_y_pixels=10) + case "RotationalMotion": + return RotationalMotionDistortion(angle_radians=torch.pi / 6) + case "SegmentedRotationalMotion": + return SegmentedRotationalMotionDistortion( + angle_radians=(0.0, torch.pi / 20, -torch.pi / 24, torch.pi / 16), + ) + case "SegmentedTranslationMotion": + return SegmentedTranslationMotionDistortion( + shift_x_pixels=(0.0, 20.0, 50.0, -50.0), + shift_y_pixels=(0.0, 10.0, -20.0, 20.0), + ) + case "GaussianKspaceBiasField": + return GaussianKspaceBiasField(width_fraction=0.35, edge_gain=0.4) + case "GaussianNoise": + return GaussianNoiseDistortion(sigma=0.00001) + case "BaseDistortion": + return BaseDistortion() + + case _: + raise ValueError(f"Unknown distortion {name!r}") diff --git a/mri_recon/examples/run_experiments_fastmri_brain.py b/mri_recon/examples/run_experiments_fastmri_brain.py new file mode 100644 index 0000000..e482e79 --- /dev/null +++ b/mri_recon/examples/run_experiments_fastmri_brain.py @@ -0,0 +1,313 @@ +from datetime import datetime +import os +import sys +import h5py +from tifffile import imwrite +import torch +import deepinv as dinv +import pandas as pd + + +sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from mri_recon.distortions import DistortedKspaceMultiCoilMRI, BaseDistortion, choose_distortion +from mri_recon.reconstruction import choose_reconstructor +from mri_recon.utils.oasis_adapter import ( + DistortedOasisMeasurement, + fastmri_measurement_to_oasis_kspace, + kspace_to_image, +) + + +sensitivity_map_estimation_algorithm = [ + "espirit", + # "unity", + # "birdcage", +] + + +DISTORTIONS = [ + "BaseDistortion", + # "PhaseEncodeGhosting", + # "CartesianUndersamplingVariableDensity", + # "CartesianUndersamplingUniformRandom", + # "HannTaperLP", + # "KaiserTaperLP", + # "RadialHighPassEmphasis", + # "IsotropicLP", + # "OffCenterAnisotropicGaussianKspaceBiasField", + # "TranslationMotion", + # "RotationalMotion", + # "SegmentedRotationalMotion", + # "SegmentedTranslationMotion", + # "GaussianKspaceBiasField", + # "GaussianNoise", +] + +RECONSTRUCTORS = [ + # "zero-filled", + # "conjugate-gradient", + # "ram", + # "dip", + # "tv-pgd", + # "wavelet-fista", + # "tv-fista", + # "tv-pdhg", + # "unet", # will trigger download of pretrained weights if not already present + # *list(EXPLICIT_UNET_ALGORITHMS) + "unet-fastmri", + "unet-oasis-acceleration4", + # "unet-oasis-acceleration8", + # "unet-oasis-acceleration10", +] + +reference_reconstructor = "conjugate-gradient" +reference_map_estimation = "espirit" + +filenames = [ + "/home/melanie.dohmen/mri_recon/data/fastmri/multicoil_brain_test/file_brain_AXFLAIR_200_6002452.h5" +] + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +result_path = "/home/melanie.dohmen/mri_recon/reports/experiments_fastmri_brain/" + +os.makedirs(result_path, exist_ok=True) + +distortion_times = {} +reconstruction_times = {} + +for f_idx, filename in enumerate(filenames): + with h5py.File(filename, "r") as hf: + kspace_data = hf["kspace"][:] + reconstruction_rss = hf["reconstruction_rss"][:] + # hdr = hf["ismrmrd_header"][()] + # print(hf.keys()) + + x = torch.from_numpy(reconstruction_rss).unsqueeze(0).unsqueeze(0) + y = torch.view_as_real(torch.from_numpy(kspace_data)).unsqueeze(0).moveaxis(-1, 1) + # image shape: (1, channels, slices, H, W) + print("image x.shape:", x.shape) + # k-space shape: (1, channels, slices, coils, H, W) + print("k-space y.shape:", y.shape) + x = x.to(device) + y = y.to(device) + + # select middle slice: + x = x[:, :, x.shape[2] // 2, ...] + y = y[:, :, y.shape[2] // 2, ...] + + # read size of dimensions: + batch_size, channels, n_coils, ksp_w, ksp_h = y.shape + + for map_estimation in sensitivity_map_estimation_algorithm: + print(f"Estimating coil maps with {map_estimation}...") + if map_estimation == "espirit": + estimate_coil_maps = dinv.datasets.MRISliceTransform( + estimate_coil_maps=True, + acs=15, # Num. low frequency, fix to 15 + ) + _, _, params = estimate_coil_maps(target=x[0], kspace=y[0]) + + coil_maps = params["coil_maps"] + print("estimated coil maps shape: ", coil_maps.shape) + + coil_maps_result_path = os.path.join( + result_path, f"brain_sample_{f_idx}_coil_maps_{map_estimation}" + ) + os.makedirs(coil_maps_result_path, exist_ok=True) + for c_idx in range(coil_maps.shape[0]): + imwrite( + os.path.join( + coil_maps_result_path, + f"brain_sample_{f_idx}_{map_estimation}_map_{c_idx}.tiff", + ), + coil_maps[c_idx].abs().numpy(), + ) + + elif map_estimation == "unity": + coil_maps = torch.ones((n_coils, ksp_w, ksp_h), dtype=torch.complex64, device=device) + print("estimated coil maps shape: ", coil_maps.shape) + coil_maps_result_path = os.path.join( + result_path, f"brain_sample_{f_idx}_coil_maps_{map_estimation}" + ) + os.makedirs(coil_maps_result_path, exist_ok=True) + for c_idx in range(coil_maps.shape[0]): + imwrite( + os.path.join( + coil_maps_result_path, + f"brain_sample_{f_idx}_{map_estimation}_map_{c_idx}.tiff", + ), + coil_maps[c_idx].abs().numpy(), + ) + + elif map_estimation == "birdcage": + coil_maps = n_coils + + if map_estimation == reference_map_estimation: + physics_clean = dinv.physics.MultiCoilMRI( + img_size=(ksp_w, ksp_h), + mask=None, + coil_maps=coil_maps, + device=device, + ) + + y_distorted = BaseDistortion()(y) + print("base distortion does not change y:", torch.all(y_distorted == y)) + print("kspace distorted.shape: ", y_distorted.shape) + + x_recon_reference = choose_reconstructor(reference_reconstructor)( + y_distorted, physics_clean + ) + print("reconstructed image shape: ", x_recon_reference.shape) + + x_recon_reference_cropped = physics_clean.crop(x_recon_reference, shape=x.shape) + print("cropped reconstructed image shape: ", x_recon_reference_cropped.shape) + + x_as_inversed_y = physics_clean.A_adjoint(y_distorted, rss=True) + print("inversed y shape: ", x_as_inversed_y.shape) + x_as_inversed_y_cropped = physics_clean.crop(x_as_inversed_y, shape=x.shape) + print("cropped inversed y shape: ", x_as_inversed_y_cropped.shape) + + imwrite( + os.path.join( + result_path, + f"brain_sample_{f_idx}_{map_estimation}_reconstructed_reference.tiff", + ), + x_recon_reference_cropped[0, 0].abs().numpy(), + ) + imwrite( + os.path.join(result_path, f"brain_sample_{f_idx}_{map_estimation}_inversed_y.tiff"), + x_as_inversed_y_cropped[0, 0].abs().numpy(), + ) + + if map_estimation == "birdcage": + coil_maps = physics_clean.coil_maps + print("estimated coil maps shape: ", coil_maps.shape) + coil_maps_result_path = os.path.join( + result_path, f"brain_sample_{f_idx}_coil_maps_{map_estimation}" + ) + os.makedirs(coil_maps_result_path, exist_ok=True) + for c_idx in range(coil_maps.shape[0]): + imwrite( + os.path.join( + coil_maps_result_path, + f"brain_sample_{f_idx}_{map_estimation}_map_{c_idx}.tiff", + ), + coil_maps[0, c_idx].abs().numpy(), + ) + + for distortion_name in DISTORTIONS: + print("Distortion: ", distortion_name) + + start = datetime.now() + + distortion = choose_distortion(distortion_name) + physics_distorted = DistortedKspaceMultiCoilMRI( + distortion=distortion, + img_size=(ksp_w, ksp_h), + mask=None, + coil_maps=coil_maps, + device=device, + ) + + print("inserting k-space into distortion with shape: ", y.shape) + y_distorted = distortion(y) + + x_distorted_as_inversed_y = physics_distorted.A_adjoint(y_distorted, rss=True) + x_distorted_as_inversed_y_cropped = physics_distorted.crop( + x_distorted_as_inversed_y, shape=x.shape + ) + imwrite( + os.path.join( + result_path, + f"brain_sample_{f_idx}_{map_estimation}_{distortion_name}_inversed_y.tiff", + ), + x_distorted_as_inversed_y_cropped[0, 0].abs().numpy(), + ) + + for reconstructor_name in RECONSTRUCTORS: + print("Reconstructor: ", reconstructor_name) + + start_recon = datetime.now() + + if reconstructor_name in [ + "unet-oasis-acceleration4", + "unet-oasis-acceleration8", + "unet-oasis-acceleration10", + ]: + physics_distorted = DistortedOasisMeasurement( + distortion=distortion, + img_size=(ksp_w, ksp_h), + mask=None, + coil_maps=coil_maps, + device=device, + ) + y_distorted_for_recon = kspace_to_image( + fastmri_measurement_to_oasis_kspace(y_distorted) + ) + else: + y_distorted_for_recon = y_distorted + + reconstructor = choose_reconstructor( + reconstructor_name, + img_size=y.shape[-2:], + device=device, + verbose=True, + ).to(device) + + try: + x_distorted = reconstructor(y_distorted_for_recon, physics_distorted) + + x_distorted_cropped = ( + physics_distorted.crop(x_distorted, shape=x.shape).detach().cpu() + ) + + imwrite( + os.path.join( + result_path, + f"brain_sample_{f_idx}_{map_estimation}_{distortion_name}_{reconstructor_name}.tiff", + ), + physics_distorted.coil_maps[:, 0].abs().numpy(), + ) + imwrite( + os.path.join( + result_path, + f"brain_sample_{f_idx}_{map_estimation}_{distortion_name}_{reconstructor_name}_reconstructed.tiff", + ), + x_distorted_cropped[0, 0].abs().numpy(), + ) + imwrite( + os.path.join( + result_path, + f"brain_sample_{f_idx}_{map_estimation}_{distortion_name}_{reconstructor_name}_reconstructed.tiff", + ), + x_distorted_cropped[0, 0].abs().numpy(), + ) + except Exception as e: + print(f"Reconstruction with {reconstructor_name} failed due to error: {e}") + continue + + end_recon = datetime.now() + if reconstructor_name not in reconstruction_times: + reconstruction_times[reconstructor_name] = [ + (end_recon - start_recon).total_seconds() + ] + else: + reconstruction_times[reconstructor_name].append( + (end_recon - start_recon).total_seconds() + ) + + end = datetime.now() + if distortion_name not in distortion_times: + distortion_times[distortion_name] = [(end - start).total_seconds()] + else: + distortion_times[distortion_name].append((end - start).total_seconds()) + + +distortion_times_df = pd.DataFrame(distortion_times, index=[sensitivity_map_estimation_algorithm]) +distortion_times_df.to_csv(os.path.join(result_path, "distortion_times.csv"), index=False) +reconstruction_times_df = pd.DataFrame(reconstruction_times) +reconstruction_times_df.to_csv(os.path.join(result_path, "reconstruction_times.csv"), index=False) + +print(distortion_times_df.mean()) +print(reconstruction_times_df.mean()) diff --git a/mri_recon/examples/run_experiments_fastmri_prostateT2.py b/mri_recon/examples/run_experiments_fastmri_prostateT2.py new file mode 100644 index 0000000..e999d56 --- /dev/null +++ b/mri_recon/examples/run_experiments_fastmri_prostateT2.py @@ -0,0 +1,289 @@ +import os +import sys +import h5py +from tifffile import imwrite +import torch +import deepinv as dinv + +import numpy as np + + +sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from mri_recon.distortions import DistortedKspaceMultiCoilMRI, BaseDistortion +from mri_recon.reconstruction import choose_reconstructor +from mri_recon.utils import Grappa + + +sensitivity_map_estimation_algorithm = [ + # "espirit", + "unity", + # "birdcage", +] + + +DISTORTIONS = [ + # "PhaseEncodeGhosting", + # "CartesianUndersamplingVariableDensityRandom", + "CartesianUndersamplingUniformRandom", + # "HannTaperLP", + # "KaiserTaperLP", + # "RadialHighPassEmphasis", + # "IsotropicLP", + # "OffCenterAnisotropicGaussianBiasField", + # "TranslationMotion", + # "RotationalMotion", + # "SegmentedRotationalMotion", + # "SegmentedTranslationMotion", + "GaussianKspaceBiasField", + # "GaussianNoise", +] + +RECONSTRUCTORS = [ + "zero-filled", + "conjugate-gradient", + # "ram", + # "dip", + # "tv-pgd", + # "wavelet-fista", + # "tv-fista", + # "tv-pdhg", + # "unet", # will trigger download of pretrained weights if not already present + # *list(EXPLICIT_UNET_ALGORITHMS) + #'unet-fastmri', + #'unet-oasis-acceleration4', + #'unet-oasis-acceleration8', + #'unet-oasis-acceleration10', +] + +reference_reconstructor = "conjugate-gradient" +reference_map_estimation = "unity" + +filenames = [ + "/home/melanie.dohmen/mri_recon/data/fastmri/fastMRI_prostate_T2_IDS_001_020/file_prostate_AXT2_001.h5" +] + + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +result_path = "/home/melanie.dohmen/mri_recon/reports/experiments_fastmri_prostateT2/" + +os.makedirs(result_path, exist_ok=True) + + +def correct_kspace_data_with_calibration( + y: torch.Tensor, calibration_data: torch.Tensor +) -> torch.Tensor: + n_avg, n_slices, n_coils, ksp_w, ksp_h = kspace_data.shape + + # Calib_data shape: num_slices, num_coils, num_pe_cal + grappa_weight_dict = {} + grappa_weight_dict_2 = {} + + kspace_slice_regridded = kspace_data[0, 0, ...] + grappa_obj = Grappa( + np.transpose(kspace_slice_regridded, (2, 0, 1)), kernel_size=(5, 5), coil_axis=1 + ) + + kspace_slice_regridded_2 = kspace_data[1, 0, ...] + grappa_obj_2 = Grappa( + np.transpose(kspace_slice_regridded_2, (2, 0, 1)), kernel_size=(5, 5), coil_axis=1 + ) + + # calculate GRAPPA weights + for slice_num in range(n_slices): + calibration_regridded = calibration_data[slice_num, ...] + grappa_weight_dict[slice_num] = grappa_obj.compute_weights( + np.transpose(calibration_regridded, (2, 0, 1)) + ) + grappa_weight_dict_2[slice_num] = grappa_obj_2.compute_weights( + np.transpose(calibration_regridded, (2, 0, 1)) + ) + + # apply GRAPPA weights + kspace_post_grappa_all = np.zeros(shape=kspace_data.shape, dtype=complex) + + for average, grappa_obj, grappa_weight_dict in zip( + [0, 1, 2], + [grappa_obj, grappa_obj_2, grappa_obj], + [grappa_weight_dict, grappa_weight_dict_2, grappa_weight_dict], + ): + for slice_num in range(n_slices): + kspace_slice_regridded = kspace_data[average, slice_num, ...] + kspace_post_grappa = grappa_obj.apply_weights( + np.transpose(kspace_slice_regridded, (2, 0, 1)), grappa_weight_dict[slice_num] + ) + kspace_post_grappa_all[average, slice_num, ...] = np.moveaxis( + np.moveaxis(kspace_post_grappa, 0, 1), 1, 2 + ) + + return kspace_post_grappa_all + + +for f_idx, filename in enumerate(filenames): + with h5py.File(filename, "r") as hf: + kspace_data = hf["kspace"][:] + calibration_data = hf["calibration_data"][:] + hdr = hf["ismrmrd_header"][()] + reconstruction_rss = hf["reconstruction_rss"][:] + atts = dict() + atts["max"] = hf.attrs["max"] + atts["norm"] = hf.attrs["norm"] + atts["patient_id"] = hf.attrs["patient_id"] + atts["acquisition"] = hf.attrs["acquisition"] + + n_avg, n_slices, n_coils, ksp_w, ksp_h = kspace_data.shape + + # correct k-space data with calibration data: + + kspace_data = correct_kspace_data_with_calibration(kspace_data, calibration_data) + + x = torch.from_numpy(reconstruction_rss).unsqueeze(0).unsqueeze(0) + y = torch.view_as_real(torch.from_numpy(kspace_data)).unsqueeze(0).moveaxis(-1, 1) + # image shape: (1, slices, H, W) + print("image x.shape:", x.shape) + # k-space shape: (1, slices, coils, H, W) + print("k-space y.shape:", y.shape) + x = x.to(device) + y = y.to(device) + + # select middle slice: + x = x[:, :, x.shape[2] // 2, ...] + y = y[:, :, y.shape[2] // 2, ...] + + recon_1 = [] + + for ave in range(n_avg): + # estimate coil maps: + for map_estimation in sensitivity_map_estimation_algorithm: + if map_estimation == "espirit": + estimate_coil_maps = dinv.datasets.MRISliceTransform( + estimate_coil_maps=True, + acs=15, # Num. low frequency, fix to 15 + ) + _, _, params = estimate_coil_maps(target=x[0], kspace=y[0]) + + coil_maps = params["coil_maps"] + print("estimated coil maps shape: ", coil_maps.shape) + imwrite( + os.path.join( + result_path, f"prostate_sample_{f_idx}_{map_estimation}_maps.tiff" + ), + coil_maps[:, 0].abs().numpy(), + ) + + elif map_estimation == "unity": + coil_maps = torch.ones( + (n_coils, ksp_w, ksp_h), dtype=torch.complex64, device=device + ) + print("estimated coil maps shape: ", coil_maps.shape) + imwrite( + os.path.join( + result_path, f"prostate_sample_{f_idx}_{map_estimation}_maps.tiff" + ), + coil_maps[:, 0].abs().numpy(), + ) + + elif map_estimation == "birdcage": + coil_maps = n_coils + + if map_estimation == reference_map_estimation: + physics_clean = dinv.physics.MultiCoilMRI( + img_size=(ksp_w, ksp_h), + mask=None, + coil_maps=coil_maps, + device=device, + ) + + y_distorted = BaseDistortion()(y) + print("base distortion does not change y:", torch.all(y_distorted == y)) + print("kspace distorted.shape: ", y_distorted.shape) + + x_recon_reference = choose_reconstructor(reference_reconstructor)( + y_distorted, physics_clean + ) + print("reconstructed image shape: ", x_recon_reference.shape) + + x_recon_reference_cropped = physics_clean.crop(x_recon_reference, shape=x.shape) + print("cropped reconstructed image shape: ", x_recon_reference_cropped.shape) + + recon_1.append(x_recon_reference_cropped[0, 0].abs().numpy()) + + x_as_inversed_y = physics_clean.A_adjoint(y_distorted, rss=True) + print("inversed y shape: ", x_as_inversed_y.shape) + x_as_inversed_y_cropped = physics_clean.crop(x_as_inversed_y, shape=x.shape) + print("cropped inversed y shape: ", x_as_inversed_y_cropped.shape) + + imwrite( + os.path.join( + result_path, + f"prostate_sample_{f_idx}_{map_estimation}_reconstructed_reference.tiff", + ), + x_recon_reference_cropped[0, 0].abs().numpy(), + ) + imwrite( + os.path.join( + result_path, f"prostate_sample_{f_idx}_{map_estimation}_inversed_y.tiff" + ), + x_as_inversed_y_cropped[0, 0].abs().numpy(), + ) + + if map_estimation == "birdcage": + coil_maps = physics_clean.coil_maps + print("estimated coil maps shape: ", coil_maps.shape) + imwrite( + os.path.join( + result_path, f"prostate_sample_{f_idx}_{map_estimation}_maps.tiff" + ), + coil_maps[:, 0].abs().numpy(), + ) + + for distortion_name in DISTORTIONS: + physics_distorted = DistortedKspaceMultiCoilMRI( + distortion=BaseDistortion(), + img_size=(ksp_w, ksp_h), + mask=None, + coil_maps=coil_maps, + device=device, + ) + + x_distorted_as_inversed_y = physics_distorted.A_adjoint(y_distorted, rss=True) + x_distorted_as_inversed_y_cropped = physics_distorted.crop( + x_distorted_as_inversed_y, shape=x.shape + ) + imwrite( + os.path.join( + result_path, f"prostate_sample_{f_idx}_{distortion_name}_inversed_y.tiff" + ), + x_distorted_as_inversed_y_cropped[0, 0].abs().numpy(), + ) + + for reconstructor_name in RECONSTRUCTORS: + reconstructor = choose_reconstructor(reconstructor_name) + + x_distorted = reconstructor(y_distorted, physics_distorted) + + x_distorted_cropped = physics_distorted.crop(x_distorted, shape=x.shape) + + imwrite( + os.path.join( + result_path, + f"prostate_sample_{f_idx}_{distortion_name}_{reconstructor_name}.tiff", + ), + physics_distorted.coil_maps[:, 0].abs().numpy(), + ) + imwrite( + os.path.join( + result_path, + f"prostate_sample_{f_idx}_{distortion_name}_{reconstructor_name}_reconstructed.tiff", + ), + x_distorted_cropped[0, 0].abs().numpy(), + ) + + recon_1_mean = np.mean(recon_1, axis=0) + imwrite( + os.path.join( + result_path, + f"prostate_sample_{f_idx}_{reference_map_estimation}_reconstructed_reference_mean.tiff", + ), + recon_1_mean, + ) diff --git a/mri_recon/reconstruction/inference.py b/mri_recon/reconstruction/inference.py index b0fe1d3..27c1c65 100644 --- a/mri_recon/reconstruction/inference.py +++ b/mri_recon/reconstruction/inference.py @@ -51,6 +51,23 @@ def validate_algorithm_dataset_compatibility(dataset: str, algorithm: str) -> No "The algorithm 'unet-fastmri' is not supported on the OASIS dataset. " "Use one of the explicit OASIS U-Net algorithms instead." ) + elif dataset == "fastmri" and algorithm in OASIS_UNET_ALGORITHMS: + raise ValueError( + "The algorithm 'unet-oasis' is not supported on the FastMRI dataset. " + "Use the 'unet-fastmri' algorithm instead." + ) + elif dataset == "fastmri-multicoil" and algorithm == FASTMRI_UNET_ALGORITHM: + raise ValueError( + "The algorithm 'unet-fastmri' (knee) is not supported on the FastMRI multicoil (brain) dataset. " + "Use the 'unet-oasis' algorithm instead." + ) + elif dataset in ["cmrxrecon", "prostate"] and algorithm in [FASTMRI_UNET_ALGORITHM] + list( + OASIS_UNET_ALGORITHMS.keys() + ): + raise ValueError( + f"The algorithm {algorithm} ({'heart' if dataset == 'cmrxrecon' else 'prostate'}) is not supported on the cmrxrecon or prostate datasets. " + "No trained unet model available for this dataset." + ) def choose_reconstructor( diff --git a/mri_recon/utils/__init__.py b/mri_recon/utils/__init__.py index 744c860..de656c6 100644 --- a/mri_recon/utils/__init__.py +++ b/mri_recon/utils/__init__.py @@ -4,13 +4,19 @@ from .io import matches_sha256 as matches_sha256 from .oasis_adapter import OasisCenteredFFTPhysics as OasisCenteredFFTPhysics from .oasis_adapter import OasisSliceDataset as OasisSliceDataset +from .oasis_adapter import OasisCenterSliceFolderDataset as OasisCenterSliceFolderDataset from .oasis_adapter import fastmri_measurement_to_image as fastmri_measurement_to_image from .oasis_adapter import ( fastmri_measurement_to_oasis_kspace as fastmri_measurement_to_oasis_kspace, ) +from .oasis_adapter import ( + oasis_kspace_to_fastmri_measurement as oasis_kspace_to_fastmri_measurement, +) from .oasis_adapter import image_to_kspace as image_to_kspace from .oasis_adapter import kspace_to_image as kspace_to_image +from .prostate_adaptor import FastMRIProstateDataset as FastMRIProstateDataset from .plot import save_kspace_plot as save_kspace_plot +from .plot import _kspace_to_log_magnitude as _kspace_to_log_magnitude __all__ = [ "download_file_with_sha256", @@ -23,5 +29,6 @@ "matches_sha256", "OasisCenteredFFTPhysics", "OasisSliceDataset", + "FastMRIProstateDataset", "save_kspace_plot", ] diff --git a/mri_recon/utils/grappa.py b/mri_recon/utils/grappa.py new file mode 100644 index 0000000..a2cbe9c --- /dev/null +++ b/mri_recon/utils/grappa.py @@ -0,0 +1,215 @@ +from typing import Dict, Tuple +import numpy as np +from skimage.util import view_as_windows +from tempfile import NamedTemporaryFile as NTF + + +class Grappa: + def __init__( + self, kspace: np.ndarray, kernel_size: Tuple[int, int] = (5, 5), coil_axis: int = -1 + ) -> None: + self.kspace = kspace + self.kernel_size = kernel_size + self.coil_axis = coil_axis + self.lamda = 0.01 + + self.kernel_var_dict = self.get_kernel_geometries() + + def get_kernel_geometries(self): + """ + Extract unique kernel geometries based on a slice of kspace data + + Returns + ------- + geometries : dict + A dictionary containing the following keys: + - 'patches': an array of overlapping patches from the k-space data. + - 'patch_indices': an array of unique patch indices. + - 'holes_x': a dictionary of x-coordinates for holes in each patch. + - 'holes_y': a dictionary of y-coordinates for holes in each patch. + + Notes + ----- + This function extracts unique kernel geometries from a slice of k-space data. + The geometries correspond to overlapping patches that contain at least one hole. + A hole is defined as a region of k-space data where the absolute value of the + complex signal is equal to zero. The function returns a dictionary containing + information about the patches and holes, which can be used to compute weights + for each geometry using the GRAPPA algorithm. + + """ + self.kspace = np.moveaxis(self.kspace, self.coil_axis, -1) + + # Quit early if there are no holes + if np.sum((np.abs(self.kspace[..., 0]) == 0).flatten()) == 0: + return np.moveaxis(self.kspace, -1, self.coil_axis) + + kx, ky = self.kernel_size[:] + kx2, ky2 = int(kx / 2), int(ky / 2) + nc = self.kspace.shape[-1] + + self.kspace = np.pad(self.kspace, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") + + mask = np.ascontiguousarray(np.abs(self.kspace[..., 0]) > 0) + + with NTF() as fP: + # Get all overlapping patches from the mask + P = np.memmap( + fP, + dtype=mask.dtype, + mode="w+", + shape=(mask.shape[0] - 2 * kx2, mask.shape[1] - 2 * ky2, 1, kx, ky), + ) + P = view_as_windows(mask, (kx, ky)) + Psh = P.shape[:] # save shape for unflattening indices later + P = P.reshape((-1, kx, ky)) + + # Find the unique patches and associate them with indices + P, iidx = np.unique(P, return_inverse=True, axis=0) + + # Filter out geometries that don't have a hole at the center. + # These are all the kernel geometries we actually need to + # compute weights for. + validP = np.argwhere(~P[:, kx2, ky2]).squeeze() + + # ignore empty patches + invalidP = np.argwhere(np.all(P == 0, axis=(1, 2))) + validP = np.setdiff1d(validP, invalidP, assume_unique=True) + + validP = np.atleast_1d(validP) + + # Give P back its coil dimension + P = np.tile(P[..., None], (1, 1, 1, nc)) + + holes_x = {} + holes_y = {} + for ii in validP: + # x, y define where top left corner is, so move to ctr, + # also make sure they are iterable by enforcing atleast_1d + idx = np.unravel_index(np.argwhere(iidx == ii), Psh[:2]) + x, y = idx[0] + kx2, idx[1] + ky2 + x = np.atleast_1d(x.squeeze()) + y = np.atleast_1d(y.squeeze()) + + holes_x[ii] = x + holes_y[ii] = y + + return {"patches": P, "patch_indices": validP, "holes_x": holes_x, "holes_y": holes_y} + + def compute_weights(self, calib: np.ndarray) -> Dict[int, np.ndarray]: + """ + Compute the GRAPPA weights for each slice in the input calibration data. + + Parameters: + ---------- + calib : numpy.ndarray + Calibration data with shape (Nx, Nc, Ny) where Nx, Ny are the size of the image in the x and y dimensions, + respectively, and Nc is the number of coils. + + Returns: + ------- + weights : dict + A dictionary of GRAPPA weights for each patch index. + + Notes: + ----- + The GRAPPA algorithm is used to estimate the missing k-space data in undersampled MRI acquisitions. + The algorithm used to compute the GRAPPA weights involves first extracting patches from the calibration data, + and then solving a linear system to estimate the weights. The resulting weights are stored in a dictionary + where the key is the patch index. The equation to solve for the weights involves taking the product of the + sources and the targets in the patch domain, and then regularizing the matrix using Tikhonov regularization. + The function uses numpy's `memmap` to store temporary files to avoid overwhelming memory usage. + """ + + calib = np.moveaxis(calib, self.coil_axis, -1) + kx, ky = self.kernel_size[:] + kx2, ky2 = int(kx / 2), int(ky / 2) + nc = calib.shape[-1] + + calib = np.pad(calib, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") + + # Store windows in temporary files so we don't overwhelm memory + with NTF() as fA: + # Get all overlapping patches of ACS + try: + A = np.memmap( + fA, + dtype=calib.dtype, + mode="w+", + shape=(calib.shape[0] - 2 * kx, calib.shape[1] - 2 * ky, 1, kx, ky, nc), + ) + A[:] = view_as_windows(calib, (kx, ky, nc)).reshape((-1, kx, ky, nc)) + except ValueError: + A = view_as_windows(calib, (kx, ky, nc)).reshape((-1, kx, ky, nc)) + + weights = {} + + for ii in self.kernel_var_dict["patch_indices"]: + # Get the sources by masking all patches of the ACS and + # get targets by taking the center of each patch. Source + # and targets will have the following sizes: + # S : (# samples, N possible patches in ACS) + # T : (# coils, N possible patches in ACS) + # Solve the equation for the weights: using numpy.linalg.solve, + # and Tikhonov regularization for better conditioning: + # SW = T + # S^HSW = S^HT + # W = (S^HS)^-1 S^HT + # -> W = (S^HS + lamda I)^-1 S^HT + + S = A[:, self.kernel_var_dict["patches"][ii, ...]] + T = A[:, kx2, ky2, :] + ShS = S.conj().T @ S + ShT = S.conj().T @ T + lamda0 = self.lamda * np.linalg.norm(ShS) / ShS.shape[0] + weights[ii] = np.linalg.solve(ShS + lamda0 * np.eye(ShS.shape[0]), ShT).T + + return weights + + def apply_weights(self, kspace: np.ndarray, weights: Dict[int, np.ndarray]) -> np.ndarray: + """ + Applies the computed GRAPPA weights to the k-space data. + + Parameters: + ---------- + kspace : numpy.ndarray + The k-space data to apply the weights to. + + weights : dict + A dictionary containing the GRAPPA weights to apply. + + Returns: + ------- + numpy.ndarray: The reconstructed data after applying the weights. + """ + + # fin_shape = kspace.shape[:] + + # Put the coil dimension at the end + kspace = np.moveaxis(kspace, self.coil_axis, -1) + + # Get shape of kernel + kx, ky = self.kernel_size[:] + kx2, ky2 = int(kx / 2), int(ky / 2) + + # adjustment factor for odd kernel size + adjx = np.mod(kx, 2) + adjy = np.mod(ky, 2) + + # Pad kspace data + kspace = np.pad(kspace, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") + + with NTF() as frecon: + # Initialize recon array + recon = np.memmap(frecon, dtype=kspace.dtype, mode="w+", shape=kspace.shape) + + for ii in self.kernel_var_dict["patch_indices"]: + for xx, yy in zip( + self.kernel_var_dict["holes_x"][ii], self.kernel_var_dict["holes_y"][ii] + ): + # Collect sources for this hole and apply weights + S = kspace[xx - kx2 : xx + kx2 + adjx, yy - ky2 : yy + ky2 + adjy, :] + S = S[self.kernel_var_dict["patches"][ii, ...]] + recon[xx, yy, :] = (weights[ii] @ S[:, None]).squeeze() + + return np.moveaxis((recon[:] + kspace)[kx2:-kx2, ky2:-ky2, :], -1, self.coil_axis) diff --git a/mri_recon/utils/oasis_adapter.py b/mri_recon/utils/oasis_adapter.py index 0f833aa..e87cd8f 100644 --- a/mri_recon/utils/oasis_adapter.py +++ b/mri_recon/utils/oasis_adapter.py @@ -7,6 +7,7 @@ import numpy as np import torch from torch.utils.data import Dataset +import deepinv as dinv from mri_recon.distortions import BaseDistortion, DistortedKspaceMultiCoilMRI @@ -122,6 +123,76 @@ def _get_volume(self, subject_id: str) -> np.ndarray: return volume +class OasisCenterSliceFolderDataset(Dataset): + """Load 2D OASIS slices from Analyze/NIfTI volumes. + Select a center slice from all subjects in the folder. + + Parameters + ---------- + data_path : Path + Root directory containing OASIS subject folders. + + """ + + def __init__( + self, + data_path: Path, + ) -> None: + try: + import nibabel as nib + except ImportError as exc: + raise ImportError( + "OASIS loading requires nibabel. Install the project dependencies " + "or add nibabel to your environment before using OasisSliceDataset." + ) from exc + + self._nib = nib + self.data_path = Path(data_path) + self.subject_paths = self._discover_subject_paths() + + def __len__(self) -> int: + """Return the number of available slices.""" + + return len(self.subject_paths) + + def __getitem__(self, index: int) -> dict[str, object]: + """Return one complex-valued OASIS slice in repo tensor convention.""" + + volume = self._get_volume(list(self.subject_paths.values())[0]) + n_slices, _, _ = volume.shape + slice_num = n_slices // 2 + subject_id = list(self.subject_paths.keys())[0] + target_np = np.ascontiguousarray(volume[slice_num], dtype=np.float32) + real = torch.from_numpy(target_np) + x = torch.stack([real, torch.zeros_like(real)], dim=0) + return {"x": x.float(), "subject_id": subject_id, "slice_num": slice_num} + + def _discover_subject_paths(self) -> dict[str, Path]: + subject_paths = {} + for subject_dir in sorted(self.data_path.iterdir()): + if not subject_dir.is_dir(): + continue + image_glob = subject_dir / "PROCESSED" / "MPRAGE" / "T88_111" + matches = sorted(image_glob.glob("*t88_gfc.img")) + if matches: + subject_paths[subject_dir.name] = matches[0] + + if not subject_paths: + raise FileNotFoundError( + "Could not find OASIS subject folders under " + f"{self.data_path} matching PROCESSED/MPRAGE/T88_111/*t88_gfc.img." + ) + return subject_paths + + def _get_volume(self, subject_path: str) -> np.ndarray: + image_data = self._nib.load(subject_path).get_fdata(dtype=np.float32) + volume = np.ascontiguousarray( + np.transpose(np.squeeze(image_data), (1, 0, 2)), + dtype=np.float32, + ) + return volume + + def image_to_kspace(x: torch.Tensor) -> torch.Tensor: """Convert channel-first complex images to centered k-space. @@ -169,6 +240,8 @@ def kspace_to_image(y: torch.Tensor) -> torch.Tensor: def fastmri_measurement_to_image( y: torch.Tensor, + coil_maps: torch.Tensor | None = None, + rss: bool = False, device: torch.device | str | None = None, ) -> torch.Tensor: """Convert FastMRI measurements to image space using the repo's native physics. @@ -177,6 +250,11 @@ def fastmri_measurement_to_image( ---------- y : torch.Tensor FastMRI measurement tensor with shape ``(B, 2, H, W)``. + coil_maps : torch.Tensor | None, optional + Coil sensitivity maps with shape ``(B, C, H, W)``, where ``C`` is the number of coils. + rss : bool, optional + If ``True``, return root-sum-of-squares image across coils. Otherwise, + return coil-combined image using the provided coil sensitivity maps. Defaults to ``False``. device : torch.device | str, optional Device on which to instantiate the temporary native physics operator. @@ -188,16 +266,17 @@ def fastmri_measurement_to_image( if device is None: device = y.device - physics = DistortedKspaceMultiCoilMRI( - distortion=BaseDistortion(), + physics = dinv.physics.MultiCoilMRI( img_size=(1, 2, *y.shape[-2:]), + coil_maps=coil_maps, device=device, ) - return physics.A_adjoint(y) + return physics.A_adjoint(y, rss=rss) def fastmri_measurement_to_oasis_kspace( y: torch.Tensor, + coil_maps: torch.Tensor | None = None, device: torch.device | str | None = None, ) -> torch.Tensor: """Adapt FastMRI measurements to the centered OASIS k-space convention. @@ -206,6 +285,8 @@ def fastmri_measurement_to_oasis_kspace( ---------- y : torch.Tensor FastMRI measurement tensor with shape ``(B, 2, H, W)``. + coil_maps : torch.Tensor | None, optional + Coil sensitivity maps with shape ``(B, C, H, W)``, where ``C`` is the number of coils. device : torch.device | str, optional Device on which to instantiate the temporary native physics operator. @@ -215,10 +296,68 @@ def fastmri_measurement_to_oasis_kspace( Centered OASIS-convention k-space tensor with shape ``(B, 2, H, W)``. """ - return image_to_kspace(fastmri_measurement_to_image(y, device=device)) + return image_to_kspace(fastmri_measurement_to_image(y, coil_maps=coil_maps, device=device)) -class OasisCenteredFFTPhysics: +def image_to_fast_mri_measurement( + x: torch.Tensor, + coil_maps: torch.Tensor | None = None, + device: torch.device | str | None = None, +) -> torch.Tensor: + """Perform FFT from image space to k-space (fast-MRI convention). + + Parameters + ---------- + y : torch.Tensor + OASIS k-space measurement tensor with shape ``(B, 2, H, W)``. + coil_maps : torch.Tensor | None, optional + Coil sensitivity maps with shape ``(B, C, H, W)``, where ``C`` is the number of coils. + device : torch.device | str, optional + Device on which to instantiate the temporary native physics operator. + + Returns + ------- + torch.Tensor + FastMRI-convention k-space tensor with shape ``(B, 2, H, W)``. + """ + + if device is None: + device = x.device + physics = DistortedKspaceMultiCoilMRI( + distortion=BaseDistortion(), + img_size=(1, 2, *x.shape[-2:]), + coil_maps=coil_maps, + device=device, + ) + return physics.A(x) + + +def oasis_kspace_to_fastmri_measurement( + y: torch.Tensor, + coil_maps: torch.Tensor | None = None, + device: torch.device | str | None = None, +) -> torch.Tensor: + """Adapt OASIS-convention k-space to FastMRI measurement convention. + + Parameters + ---------- + y : torch.Tensor + Centered OASIS-convention k-space tensor with shape ``(B, 2, H, W)``. + coil_maps : torch.Tensor | None, optional + Coil sensitivity maps with shape ``(B, C, H, W)``, where ``C`` is the number of coils. + device : torch.device | str, optional + Device on which to instantiate the temporary native physics operator. + + Returns + ------- + torch.Tensor + FastMRI-convention k-space tensor with shape ``(B, 2, H, W)``. + """ + + return image_to_fast_mri_measurement(kspace_to_image(y), coil_maps=coil_maps, device=device) + + +class OasisCenteredFFTPhysics(dinv.physics.LinearPhysics): """Physics adapter matching the OASIS U-Net FFT convention. Parameters @@ -227,7 +366,8 @@ class OasisCenteredFFTPhysics: K-space distortion applied after the centered FFT. """ - def __init__(self, distortion: BaseDistortion) -> None: + def __init__(self, distortion: BaseDistortion, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) self.distortion = distortion def A(self, x: torch.Tensor) -> torch.Tensor: @@ -261,3 +401,14 @@ def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: """ return kspace_to_image(self.distortion.A_adjoint(y)) + + def A_dagger(self, y: torch.Tensor, **kwargs) -> torch.Tensor: + r""" + Computes least squares solution to the MRI inverse problem, as proposed in `SENSE: Sensitivity encoding for fast MRI `_. + + By default uses conjugate gradient solver. Overwrite default solver arguments by passing `kwargs`. See :func:`deepinv.optim.linear.least_squares` for details. + + :param dict kwargs: kwargs to pass to base :meth:`deepinv.physics.LinearPhysics.A_dagger`. + :returns: (:class:`torch.Tensor`) image with shape `(B,2,...,H,W)` + """ + return super().A_dagger(y, **kwargs) diff --git a/mri_recon/utils/plot.py b/mri_recon/utils/plot.py index 068c222..b3ff225 100644 --- a/mri_recon/utils/plot.py +++ b/mri_recon/utils/plot.py @@ -11,11 +11,19 @@ def _kspace_to_log_magnitude(kspace: torch.Tensor) -> torch.Tensor: """Convert k-space tensor to a log-magnitude image for visualization.""" - if kspace.ndim == 4: + if kspace.ndim == 5: + # show only middle coil for visualization + # (1, 2, C, H, W) -> (2, H, W) + kspace = kspace[0, :, kspace.shape[2] // 2] + elif kspace.ndim == 4: + # (1, 2, H, W) -> (2, H, W) kspace = kspace[0] - if kspace.ndim != 3 or kspace.shape[0] != 2: + elif kspace.ndim == 3: + pass + # (2, H, W) -> (2, H, W) + else: raise ValueError( - f"Expected k-space with shape (2, H, W) or (1, 2, H, W), got {tuple(kspace.shape)}" + f"Expected k-space with shape (2, H, W) or (1, 2, H, W) or (1, 2, C, H, W),got {tuple(kspace.shape)}" ) kspace = kspace.detach().cpu() @@ -43,11 +51,18 @@ def save_kspace_plot( ) -> None: """Save side-by-side log-magnitude visualizations of clean and distorted k-space.""" + print("transforming k-space to log-magnitude images for visualization...") + print(f"\tclean k-space shape: {clean_kspace.shape}") + print(f"\tdistorted k-space shape: {distorted_kspace.shape}") + images = [ ("Original k-space", _kspace_to_log_magnitude(clean_kspace)), ("Distorted k-space", _kspace_to_log_magnitude(distorted_kspace)), ] + print(f"clean k-space magnitude shape: {images[0][1].shape}") + print(f"distorted k-space magnitude shape: {images[1][1].shape}") + fig, axes = plt.subplots(1, 2, figsize=(8, 4), constrained_layout=True) fig.suptitle(f"Distortion: {distortion_label}") for ax, (title, image) in zip(axes, images, strict=True): diff --git a/mri_recon/utils/prostate_adaptor.py b/mri_recon/utils/prostate_adaptor.py new file mode 100644 index 0000000..e12f5b8 --- /dev/null +++ b/mri_recon/utils/prostate_adaptor.py @@ -0,0 +1,461 @@ +import os +import glob +from typing import Dict, Tuple, Sequence +import h5py +import numpy as np +from skimage.util import view_as_windows +from tempfile import NamedTemporaryFile as NTF +from tifffile import imwrite +import torch +import xml.etree.ElementTree as etree + + +class Grappa: + def __init__( + self, kspace: np.ndarray, kernel_size: Tuple[int, int] = (5, 5), coil_axis: int = -1 + ) -> None: + self.kspace = kspace + self.kernel_size = kernel_size + self.coil_axis = coil_axis + self.lamda = 0.01 + + self.kernel_var_dict = self.get_kernel_geometries() + + def get_kernel_geometries(self): + """ + Extract unique kernel geometries based on a slice of kspace data + + Returns + ------- + geometries : dict + A dictionary containing the following keys: + - 'patches': an array of overlapping patches from the k-space data. + - 'patch_indices': an array of unique patch indices. + - 'holes_x': a dictionary of x-coordinates for holes in each patch. + - 'holes_y': a dictionary of y-coordinates for holes in each patch. + + Notes + ----- + This function extracts unique kernel geometries from a slice of k-space data. + The geometries correspond to overlapping patches that contain at least one hole. + A hole is defined as a region of k-space data where the absolute value of the + complex signal is equal to zero. The function returns a dictionary containing + information about the patches and holes, which can be used to compute weights + for each geometry using the GRAPPA algorithm. + + """ + self.kspace = np.moveaxis(self.kspace, self.coil_axis, -1) + + # Quit early if there are no holes + if np.sum((np.abs(self.kspace[..., 0]) == 0).flatten()) == 0: + return np.moveaxis(self.kspace, -1, self.coil_axis) + + kx, ky = self.kernel_size[:] + kx2, ky2 = int(kx / 2), int(ky / 2) + nc = self.kspace.shape[-1] + + self.kspace = np.pad(self.kspace, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") + + mask = np.ascontiguousarray(np.abs(self.kspace[..., 0]) > 0) + + with NTF() as fP: + # Get all overlapping patches from the mask + P = np.memmap( + fP, + dtype=mask.dtype, + mode="w+", + shape=(mask.shape[0] - 2 * kx2, mask.shape[1] - 2 * ky2, 1, kx, ky), + ) + P = view_as_windows(mask, (kx, ky)) + Psh = P.shape[:] # save shape for unflattening indices later + P = P.reshape((-1, kx, ky)) + + # Find the unique patches and associate them with indices + P, iidx = np.unique(P, return_inverse=True, axis=0) + + # Filter out geometries that don't have a hole at the center. + # These are all the kernel geometries we actually need to + # compute weights for. + validP = np.argwhere(~P[:, kx2, ky2]).squeeze() + + # ignore empty patches + invalidP = np.argwhere(np.all(P == 0, axis=(1, 2))) + validP = np.setdiff1d(validP, invalidP, assume_unique=True) + + validP = np.atleast_1d(validP) + + # Give P back its coil dimension + P = np.tile(P[..., None], (1, 1, 1, nc)) + + holes_x = {} + holes_y = {} + for ii in validP: + # x, y define where top left corner is, so move to ctr, + # also make sure they are iterable by enforcing atleast_1d + idx = np.unravel_index(np.argwhere(iidx == ii), Psh[:2]) + x, y = idx[0] + kx2, idx[1] + ky2 + x = np.atleast_1d(x.squeeze()) + y = np.atleast_1d(y.squeeze()) + + holes_x[ii] = x + holes_y[ii] = y + + return {"patches": P, "patch_indices": validP, "holes_x": holes_x, "holes_y": holes_y} + + def compute_weights(self, calib: np.ndarray) -> Dict[int, np.ndarray]: + """ + Compute the GRAPPA weights for each slice in the input calibration data. + + Parameters: + ---------- + calib : numpy.ndarray + Calibration data with shape (Nx, Nc, Ny) where Nx, Ny are the size of the image in the x and y dimensions, + respectively, and Nc is the number of coils. + + Returns: + ------- + weights : dict + A dictionary of GRAPPA weights for each patch index. + + Notes: + ----- + The GRAPPA algorithm is used to estimate the missing k-space data in undersampled MRI acquisitions. + The algorithm used to compute the GRAPPA weights involves first extracting patches from the calibration data, + and then solving a linear system to estimate the weights. The resulting weights are stored in a dictionary + where the key is the patch index. The equation to solve for the weights involves taking the product of the + sources and the targets in the patch domain, and then regularizing the matrix using Tikhonov regularization. + The function uses numpy's `memmap` to store temporary files to avoid overwhelming memory usage. + """ + + calib = np.moveaxis(calib, self.coil_axis, -1) + kx, ky = self.kernel_size[:] + kx2, ky2 = int(kx / 2), int(ky / 2) + nc = calib.shape[-1] + + calib = np.pad(calib, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") + + # Store windows in temporary files so we don't overwhelm memory + with NTF() as fA: + # Get all overlapping patches of ACS + try: + A = np.memmap( + fA, + dtype=calib.dtype, + mode="w+", + shape=(calib.shape[0] - 2 * kx, calib.shape[1] - 2 * ky, 1, kx, ky, nc), + ) + A[:] = view_as_windows(calib, (kx, ky, nc)).reshape((-1, kx, ky, nc)) + except ValueError: + A = view_as_windows(calib, (kx, ky, nc)).reshape((-1, kx, ky, nc)) + + weights = {} + + for ii in self.kernel_var_dict["patch_indices"]: + # Get the sources by masking all patches of the ACS and + # get targets by taking the center of each patch. Source + # and targets will have the following sizes: + # S : (# samples, N possible patches in ACS) + # T : (# coils, N possible patches in ACS) + # Solve the equation for the weights: using numpy.linalg.solve, + # and Tikhonov regularization for better conditioning: + # SW = T + # S^HSW = S^HT + # W = (S^HS)^-1 S^HT + # -> W = (S^HS + lamda I)^-1 S^HT + + S = A[:, self.kernel_var_dict["patches"][ii, ...]] + T = A[:, kx2, ky2, :] + ShS = S.conj().T @ S + ShT = S.conj().T @ T + lamda0 = self.lamda * np.linalg.norm(ShS) / ShS.shape[0] + weights[ii] = np.linalg.solve(ShS + lamda0 * np.eye(ShS.shape[0]), ShT).T + + return weights + + def apply_weights(self, kspace: np.ndarray, weights: Dict[int, np.ndarray]) -> np.ndarray: + """ + Applies the computed GRAPPA weights to the k-space data. + + Parameters: + ---------- + kspace : numpy.ndarray + The k-space data to apply the weights to. + + weights : dict + A dictionary containing the GRAPPA weights to apply. + + Returns: + ------- + numpy.ndarray: The reconstructed data after applying the weights. + """ + + # fin_shape = kspace.shape[:] + + # Put the coil dimension at the end + kspace = np.moveaxis(kspace, self.coil_axis, -1) + + # Get shape of kernel + kx, ky = self.kernel_size[:] + kx2, ky2 = int(kx / 2), int(ky / 2) + + # adjustment factor for odd kernel size + adjx = np.mod(kx, 2) + adjy = np.mod(ky, 2) + + # Pad kspace data + kspace = np.pad(kspace, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") + + with NTF() as frecon: + # Initialize recon array + recon = np.memmap(frecon, dtype=kspace.dtype, mode="w+", shape=kspace.shape) + + for ii in self.kernel_var_dict["patch_indices"]: + for xx, yy in zip( + self.kernel_var_dict["holes_x"][ii], self.kernel_var_dict["holes_y"][ii] + ): + # Collect sources for this hole and apply weights + S = kspace[xx - kx2 : xx + kx2 + adjx, yy - ky2 : yy + ky2 + adjy, :] + S = S[self.kernel_var_dict["patches"][ii, ...]] + recon[xx, yy, :] = (weights[ii] @ S[:, None]).squeeze() + + return np.moveaxis((recon[:] + kspace)[kx2:-kx2, ky2:-ky2, :], -1, self.coil_axis) + + +def et_query( + root: etree.Element, qlist: Sequence[str], namespace: str = "http://www.ismrm.org/ISMRMRD" +) -> str: + """ + ElementTree query function. + + This function queries an XML document using ElementTree. + + Parameters: + ----------- + root : Element + Root of the XML document to search through. + qlist : Sequence of str + A sequence of strings for nested searches, e.g., ["Encoding", "matrixSize"]. + namespace : str, optional + XML namespace to prepend query. + + Returns: + -------- + str + The retrieved data as a string. + """ + s = "." + prefix = "ismrmrd_namespace" + + ns = {prefix: namespace} + + for el in qlist: + s = s + f"//{prefix}:{el}" + + value = root.find(s, ns) + if value is None: + raise RuntimeError("Element not found") + + return str(value.text) + + +def get_padding(hdr: str) -> float: + """ + Extract the padding value from an XML header string. + + Parameters: + ----------- + hdr : str + The XML header string. + + Returns: + -------- + float + The padding value calculated as (x - max_enc)/2, where x is the readout dimension and + max_enc is the maximum phase-encoding dimension. + """ + et_root = etree.fromstring(hdr) + lims = ["encoding", "encodingLimits", "kspace_encoding_step_1"] + enc_limits_max = int(et_query(et_root, lims + ["maximum"])) + 1 + enc = ["encoding", "encodedSpace", "matrixSize"] + enc_x = int(et_query(et_root, enc + ["x"])) + padding = (enc_x - enc_limits_max) / 2 + + print(f"Padding from header: {padding}") + return padding + + +def get_padding_from_image_recon(k_space_shape: np.ndarray, image_recon_shape: np.ndarray) -> float: + """ + Calculate the padding value based on the shapes of the k-space data and the reconstructed image. + + Parameters: + ----------- + k_space_shape : np.ndarray + The shape of the k-space data, typically in the format (num_avg, num_slices, num_coils, num_ro, num_pe). + image_recon_shape : np.ndarray + The shape of the reconstructed image, typically in the format (num_avg, num_slices, num_coils, num_x, num_y). + + Returns: + -------- + float + The padding value calculated as (x - max_enc)/2, where x is the readout dimension from the k-space shape and + max_enc is the maximum phase-encoding dimension from the reconstructed image shape. + """ + enc_limits_max = image_recon_shape[ + -1 + ] # Assuming last dimension corresponds to phase-encoding direction + enc_x = k_space_shape[-2] # Assuming second to last dimension corresponds to readout direction + padding = (enc_x - enc_limits_max) / 2 + + print(f"Padding from image recon shapes: {padding}") + return padding + + +def zero_pad_kspace_slice_hdr( + hdr: str, unpadded_kspace: np.ndarray, image_recon_shape: np.ndarray +) -> np.ndarray: + """ + Perform zero-padding on k-space data to have the same number of + points in the x- and y-directions. + + Parameters + ---------- + hdr : str + The XML header string. + unpadded_kspace : array-like of shape (ro , coils, pe) + The k-space data to be padded. + + Returns + ------- + padded_kspace : ndarray of shape (ro_padded, coils, pe_padded) + The zero-padded k-space data, where ro_padded and pe_padded are + the dimensions of the readout and phase-encoding directions after + padding. + + Notes + ----- + The padding value is calculated using the `get_padding` function, which + extracts the padding value from the XML header string. If the difference + between the readout dimension and the maximum phase-encoding dimension + is not divisible by 2, the padding is applied asymmetrically, with one + side having an additional zero-padding. + + """ + padding = get_padding(hdr) + print(f"Calculated padding: {padding}") + padding2 = get_padding_from_image_recon(unpadded_kspace.shape, image_recon_shape) + print(f"Calculated padding from image recon shapes ({image_recon_shape}): {padding2}") + if padding % 2 != 0: + padding_left = int(np.floor(padding)) + padding_right = int(np.ceil(padding)) + else: + padding_left = int(padding) + padding_right = int(padding) + padded_kspace = np.pad(unpadded_kspace, ((0, 0), (0, 0), (padding_left, padding_right))) + + return padded_kspace + + +class FastMRIProstateDataset(torch.utils.data.Dataset): + def __init__(self, data_path: str, num_samples: None | int = None) -> None: + self.data_path = data_path + self.num_samples = num_samples + self.kspace_data, self.image_data = self.pre_calc_kspace_with_grappa() + + if num_samples is not None: + self.kspace_data = self.kspace_data[:num_samples] + self.image_data = self.image_data[:num_samples] + else: + self.num_samples = len(self.kspace_data) + + def pre_calc_kspace_with_grappa(self) -> np.ndarray: + kspace_result_list = [] + image_result_list = [] + for sample_idx, filename in enumerate(glob.glob(os.path.join(self.data_path, "*.h5"))): + if (self.num_samples is not None) and (sample_idx >= self.num_samples): + break + try: + with h5py.File(filename, "r") as hf: + kspace_data = hf["kspace"][:] + calibration_data = hf["calibration_data"][:] + hdr = hf["ismrmrd_header"][()] + image_recon = hf["reconstruction_rss"][:] + + except Exception as e: + print(f"Error processing file {filename}: {e}") + continue + + # kspace data: (num_avg, num_slices, num_coils, num_ro, num_pe) + # Calib_data: (num_slices, num_coils, num_pe_cal) + + # middle slice: + kspace_middle_slice = kspace_data[:, kspace_data.shape[1] // 2] + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_middle_slice.tiff", + kspace_middle_slice, + ) + cal_middle_slice = calibration_data[calibration_data.shape[0] // 2] + print( + f"Processing file {filename} with kspace shape {kspace_middle_slice.shape} and calib shape {cal_middle_slice.shape}" + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/calibration_middle_slice.tiff", + cal_middle_slice, + ) + + ####### + + kspace_slice_regridded = kspace_data[0, 0] + grappa_obj = Grappa( + np.transpose(kspace_slice_regridded, (2, 0, 1)), kernel_size=(5, 5), coil_axis=1 + ) + + kspace_slice_regridded_2 = kspace_data[1, 0] + grappa_obj_2 = Grappa( + np.transpose(kspace_slice_regridded_2, (2, 0, 1)), kernel_size=(5, 5), coil_axis=1 + ) + + # calculate GRAPPA weights for middle slice: + + calibration_regridded = calibration_data[calibration_data.shape[0] // 2, ...] + grappa_weight = grappa_obj.compute_weights( + np.transpose(calibration_regridded, (2, 0, 1)) + ) + grappa_weight2 = grappa_obj_2.compute_weights( + np.transpose(calibration_regridded, (2, 0, 1)) + ) + + #### + kspace_post_grappa_slice = np.zeros(shape=kspace_middle_slice.shape, dtype=complex) + kspace_post_grappa_slice_padded: Dict[int, np.ndarray] = {} + for average, grappa_obj, grappa_weight_dict in zip( + [0, 1, 2], + [grappa_obj, grappa_obj_2, grappa_obj], + [grappa_weight, grappa_weight2, grappa_weight], + ): + kspace_slice_regridded = kspace_middle_slice[average, ...] + kspace_post_grappa = grappa_obj.apply_weights( + np.transpose(kspace_slice_regridded, (2, 0, 1)), grappa_weight_dict + ) + kspace_post_grappa_slice[average] = np.moveaxis( + np.moveaxis(kspace_post_grappa, 0, 1), 1, 2 + ) + + # pad: + kspace_post_grappa_slice_padded[average] = zero_pad_kspace_slice_hdr( + hdr, kspace_post_grappa_slice[average], image_recon.shape + ) + + # stack k-space data for all averages: + kspace_result_list.append( + np.stack(list(kspace_post_grappa_slice_padded.values()), axis=0) + ) + image_result_list.append(image_recon) + return kspace_result_list, image_result_list + + def __len__(self) -> int: + return len(self.kspace_data) + + def __getitem__(self, idx: int) -> torch.Tensor: + kspace = self.kspace_data[idx] + return torch.from_numpy(kspace), torch.from_numpy(self.image_data[idx]) diff --git a/mri_recon/utils/recon_prostate_T2.py b/mri_recon/utils/recon_prostate_T2.py new file mode 100644 index 0000000..f90ace7 --- /dev/null +++ b/mri_recon/utils/recon_prostate_T2.py @@ -0,0 +1,700 @@ +import os +from tempfile import NamedTemporaryFile as NTF +from typing import Dict, Tuple, Optional, Sequence +import xml.etree.ElementTree as etree + +import h5py +import numpy as np +from numpy.fft import fftshift, ifftshift, ifftn +from tifffile import imwrite +from skimage.util import view_as_windows + + +def center_crop_im(im_3d: np.ndarray, crop_to_size: Tuple[int, int]) -> np.ndarray: + """ + Center crop an image to a given size. + + Parameters: + ----------- + im_3d : numpy.ndarray + Input image of shape (slices, x, y). + crop_to_size : tuple + Tuple containing the target size for x and y dimensions. + + Returns: + -------- + numpy.ndarray + Center cropped image of size {slices, x_cropped, y_cropped}. + """ + x_crop = im_3d.shape[-1] / 2 - crop_to_size[0] / 2 + y_crop = im_3d.shape[-2] / 2 - crop_to_size[1] / 2 + + return im_3d[ + :, int(y_crop) : int(crop_to_size[1] + y_crop), int(x_crop) : int(crop_to_size[0] + x_crop) + ] + + +def ifftnd(kspace: np.ndarray, axes: Optional[Sequence[int]] = [-1]) -> np.ndarray: + """ + Compute the n-dimensional inverse Fourier transform of the k-space data along the specified axes. + + Parameters: + ----------- + kspace: np.ndarray + The input k-space data. + axes: list or tuple, optional + The list of axes along which to compute the inverse Fourier transform. Default is [-1]. + + Returns: + -------- + img: ndarray + The output image after inverse Fourier transform. + """ + + if axes is None: + axes = range(kspace.ndim) + img = fftshift(ifftn(ifftshift(kspace, axes=axes), axes=axes), axes=axes) + img *= np.sqrt(np.prod(np.take(img.shape, axes))) + + return img + + +def create_coil_combined_im(multicoil_multislice_kspace: np.ndarray) -> np.ndarray: + """ + Create a coil combined image from a multicoil-multislice k-space array. + + Parameters: + ----------- + multicoil_multislice_kspace : array-like + Input k-space data with shape (slices, coils, readout, phase encode). + + Returns: + -------- + image_mat : array-like + Coil combined image data with shape (slices, x, y). + """ + + k = multicoil_multislice_kspace + image_mat = np.zeros((k.shape[0], k.shape[2], k.shape[3])) + for i in range(image_mat.shape[0]): + data_sl = k[i, :, :, :] + image = ifftnd(data_sl, [1, 2]) + image_rss = rss(image, axis=0) + image_mat[i, :, :] = np.flipud(image_rss) + if i == 15: + print("data_sl.shape:", data_sl.shape) + print("image.shape:", image.shape) + print("image flipped: ", np.flipud(image).shape) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/image_slice_15_kspace.tiff", + np.abs(data_sl[0, :, :]), + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/image_slice_15_ifft.tiff", + np.abs(image), + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/image_slice_15_rss.tiff", + image_rss, + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/image_slice_15.tiff", + image_mat[i, :, :], + ) + return image_mat + + +def rss(sig: np.ndarray, axis: int = -1) -> np.ndarray: + """ + Compute the Root Sum-of-Squares (RSS) value of a complex signal along a specified axis. + + Parameters + ---------- + sig : np.ndarray + The complex signal to compute the RMS value of. + axis : int, optional + The axis along which to compute the RMS value. Default is -1. + + Returns + ------- + rss : np.ndarray + The RSS value of the complex signal along the specified axis. + """ + return np.sqrt(np.sum(abs(sig) ** 2, axis)) + + +def et_query( + root: etree.Element, qlist: Sequence[str], namespace: str = "http://www.ismrm.org/ISMRMRD" +) -> str: + """ + ElementTree query function. + + This function queries an XML document using ElementTree. + + Parameters: + ----------- + root : Element + Root of the XML document to search through. + qlist : Sequence of str + A sequence of strings for nested searches, e.g., ["Encoding", "matrixSize"]. + namespace : str, optional + XML namespace to prepend query. + + Returns: + -------- + str + The retrieved data as a string. + """ + s = "." + prefix = "ismrmrd_namespace" + + ns = {prefix: namespace} + + for el in qlist: + s = s + f"//{prefix}:{el}" + + value = root.find(s, ns) + if value is None: + raise RuntimeError("Element not found") + + return str(value.text) + + +def get_padding(hdr: str) -> float: + """ + Extract the padding value from an XML header string. + + Parameters: + ----------- + hdr : str + The XML header string. + + Returns: + -------- + float + The padding value calculated as (x - max_enc)/2, where x is the readout dimension and + max_enc is the maximum phase-encoding dimension. + """ + et_root = etree.fromstring(hdr) + lims = ["encoding", "encodingLimits", "kspace_encoding_step_1"] + enc_limits_max = int(et_query(et_root, lims + ["maximum"])) + 1 + enc = ["encoding", "encodedSpace", "matrixSize"] + enc_x = int(et_query(et_root, enc + ["x"])) + padding = (enc_x - enc_limits_max) / 2 + + return padding + + +def zero_pad_kspace_hdr(hdr: str, unpadded_kspace: np.ndarray) -> np.ndarray: + """ + Perform zero-padding on k-space data to have the same number of + points in the x- and y-directions. + + Parameters + ---------- + hdr : str + The XML header string. + unpadded_kspace : array-like of shape (sl, ro , coils, pe) + The k-space data to be padded. + + Returns + ------- + padded_kspace : ndarray of shape (sl, ro_padded, coils, pe_padded) + The zero-padded k-space data, where ro_padded and pe_padded are + the dimensions of the readout and phase-encoding directions after + padding. + + Notes + ----- + The padding value is calculated using the `get_padding` function, which + extracts the padding value from the XML header string. If the difference + between the readout dimension and the maximum phase-encoding dimension + is not divisible by 2, the padding is applied asymmetrically, with one + side having an additional zero-padding. + + """ + padding = get_padding(hdr) + if padding % 2 != 0: + padding_left = int(np.floor(padding)) + padding_right = int(np.ceil(padding)) + else: + padding_left = int(padding) + padding_right = int(padding) + padded_kspace = np.pad(unpadded_kspace, ((0, 0), (0, 0), (0, 0), (padding_left, padding_right))) + + return padded_kspace + + +def zero_pad_kspace_slice_hdr(hdr: str, unpadded_kspace: np.ndarray) -> np.ndarray: + """ + Perform zero-padding on k-space data to have the same number of + points in the x- and y-directions. + + Parameters + ---------- + hdr : str + The XML header string. + unpadded_kspace : array-like of shape (ro , coils, pe) + The k-space data to be padded. + + Returns + ------- + padded_kspace : ndarray of shape (ro_padded, coils, pe_padded) + The zero-padded k-space data, where ro_padded and pe_padded are + the dimensions of the readout and phase-encoding directions after + padding. + + Notes + ----- + The padding value is calculated using the `get_padding` function, which + extracts the padding value from the XML header string. If the difference + between the readout dimension and the maximum phase-encoding dimension + is not divisible by 2, the padding is applied asymmetrically, with one + side having an additional zero-padding. + + """ + padding = get_padding(hdr) + if padding % 2 != 0: + padding_left = int(np.floor(padding)) + padding_right = int(np.ceil(padding)) + else: + padding_left = int(padding) + padding_right = int(padding) + padded_kspace = np.pad(unpadded_kspace, ((0, 0), (0, 0), (padding_left, padding_right))) + + return padded_kspace + + +class Grappa: + def __init__( + self, kspace: np.ndarray, kernel_size: Tuple[int, int] = (5, 5), coil_axis: int = -1 + ) -> None: + self.kspace = kspace + self.kernel_size = kernel_size + self.coil_axis = coil_axis + self.lamda = 0.01 + + self.kernel_var_dict = self.get_kernel_geometries() + + def get_kernel_geometries(self): + """ + Extract unique kernel geometries based on a slice of kspace data + + Returns + ------- + geometries : dict + A dictionary containing the following keys: + - 'patches': an array of overlapping patches from the k-space data. + - 'patch_indices': an array of unique patch indices. + - 'holes_x': a dictionary of x-coordinates for holes in each patch. + - 'holes_y': a dictionary of y-coordinates for holes in each patch. + + Notes + ----- + This function extracts unique kernel geometries from a slice of k-space data. + The geometries correspond to overlapping patches that contain at least one hole. + A hole is defined as a region of k-space data where the absolute value of the + complex signal is equal to zero. The function returns a dictionary containing + information about the patches and holes, which can be used to compute weights + for each geometry using the GRAPPA algorithm. + + """ + self.kspace = np.moveaxis(self.kspace, self.coil_axis, -1) + + # Quit early if there are no holes + if np.sum((np.abs(self.kspace[..., 0]) == 0).flatten()) == 0: + return np.moveaxis(self.kspace, -1, self.coil_axis) + + kx, ky = self.kernel_size[:] + kx2, ky2 = int(kx / 2), int(ky / 2) + nc = self.kspace.shape[-1] + + self.kspace = np.pad(self.kspace, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") + + mask = np.ascontiguousarray(np.abs(self.kspace[..., 0]) > 0) + + with NTF() as fP: + # Get all overlapping patches from the mask + P = np.memmap( + fP, + dtype=mask.dtype, + mode="w+", + shape=(mask.shape[0] - 2 * kx2, mask.shape[1] - 2 * ky2, 1, kx, ky), + ) + P = view_as_windows(mask, (kx, ky)) + Psh = P.shape[:] # save shape for unflattening indices later + P = P.reshape((-1, kx, ky)) + + # Find the unique patches and associate them with indices + P, iidx = np.unique(P, return_inverse=True, axis=0) + + # Filter out geometries that don't have a hole at the center. + # These are all the kernel geometries we actually need to + # compute weights for. + validP = np.argwhere(~P[:, kx2, ky2]).squeeze() + + # ignore empty patches + invalidP = np.argwhere(np.all(P == 0, axis=(1, 2))) + validP = np.setdiff1d(validP, invalidP, assume_unique=True) + + validP = np.atleast_1d(validP) + + # Give P back its coil dimension + P = np.tile(P[..., None], (1, 1, 1, nc)) + + holes_x = {} + holes_y = {} + for ii in validP: + # x, y define where top left corner is, so move to ctr, + # also make sure they are iterable by enforcing atleast_1d + idx = np.unravel_index(np.argwhere(iidx == ii), Psh[:2]) + x, y = idx[0] + kx2, idx[1] + ky2 + x = np.atleast_1d(x.squeeze()) + y = np.atleast_1d(y.squeeze()) + + holes_x[ii] = x + holes_y[ii] = y + + return {"patches": P, "patch_indices": validP, "holes_x": holes_x, "holes_y": holes_y} + + def compute_weights(self, calib: np.ndarray) -> Dict[int, np.ndarray]: + """ + Compute the GRAPPA weights for each slice in the input calibration data. + + Parameters: + ---------- + calib : numpy.ndarray + Calibration data with shape (Nx, Nc, Ny) where Nx, Ny are the size of the image in the x and y dimensions, + respectively, and Nc is the number of coils. + + Returns: + ------- + weights : dict + A dictionary of GRAPPA weights for each patch index. + + Notes: + ----- + The GRAPPA algorithm is used to estimate the missing k-space data in undersampled MRI acquisitions. + The algorithm used to compute the GRAPPA weights involves first extracting patches from the calibration data, + and then solving a linear system to estimate the weights. The resulting weights are stored in a dictionary + where the key is the patch index. The equation to solve for the weights involves taking the product of the + sources and the targets in the patch domain, and then regularizing the matrix using Tikhonov regularization. + The function uses numpy's `memmap` to store temporary files to avoid overwhelming memory usage. + """ + + calib = np.moveaxis(calib, self.coil_axis, -1) + kx, ky = self.kernel_size[:] + kx2, ky2 = int(kx / 2), int(ky / 2) + nc = calib.shape[-1] + + calib = np.pad(calib, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") + + # Store windows in temporary files so we don't overwhelm memory + with NTF() as fA: + # Get all overlapping patches of ACS + try: + A = np.memmap( + fA, + dtype=calib.dtype, + mode="w+", + shape=(calib.shape[0] - 2 * kx, calib.shape[1] - 2 * ky, 1, kx, ky, nc), + ) + A[:] = view_as_windows(calib, (kx, ky, nc)).reshape((-1, kx, ky, nc)) + except ValueError: + A = view_as_windows(calib, (kx, ky, nc)).reshape((-1, kx, ky, nc)) + + weights = {} + + for ii in self.kernel_var_dict["patch_indices"]: + # Get the sources by masking all patches of the ACS and + # get targets by taking the center of each patch. Source + # and targets will have the following sizes: + # S : (# samples, N possible patches in ACS) + # T : (# coils, N possible patches in ACS) + # Solve the equation for the weights: using numpy.linalg.solve, + # and Tikhonov regularization for better conditioning: + # SW = T + # S^HSW = S^HT + # W = (S^HS)^-1 S^HT + # -> W = (S^HS + lamda I)^-1 S^HT + + S = A[:, self.kernel_var_dict["patches"][ii, ...]] + T = A[:, kx2, ky2, :] + ShS = S.conj().T @ S + ShT = S.conj().T @ T + lamda0 = self.lamda * np.linalg.norm(ShS) / ShS.shape[0] + weights[ii] = np.linalg.solve(ShS + lamda0 * np.eye(ShS.shape[0]), ShT).T + + return weights + + def apply_weights(self, kspace: np.ndarray, weights: Dict[int, np.ndarray]) -> np.ndarray: + """ + Applies the computed GRAPPA weights to the k-space data. + + Parameters: + ---------- + kspace : numpy.ndarray + The k-space data to apply the weights to. + + weights : dict + A dictionary containing the GRAPPA weights to apply. + + Returns: + ------- + numpy.ndarray: The reconstructed data after applying the weights. + """ + + # fin_shape = kspace.shape[:] + + # Put the coil dimension at the end + kspace = np.moveaxis(kspace, self.coil_axis, -1) + + # Get shape of kernel + kx, ky = self.kernel_size[:] + kx2, ky2 = int(kx / 2), int(ky / 2) + + # adjustment factor for odd kernel size + adjx = np.mod(kx, 2) + adjy = np.mod(ky, 2) + + # Pad kspace data + kspace = np.pad(kspace, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") + + with NTF() as frecon: + # Initialize recon array + recon = np.memmap(frecon, dtype=kspace.dtype, mode="w+", shape=kspace.shape) + map_of_holes = np.zeros(shape=kspace.shape[:2], dtype=bool) + for patch_index, ii in enumerate(self.kernel_var_dict["patch_indices"]): + imwrite( + f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/patch_{ii}.tiff", + self.kernel_var_dict["patches"][ii, ...], + ) + map_of_holes = np.zeros(shape=kspace.shape, dtype=bool) + for hole_idx, (xx, yy) in enumerate( + zip(self.kernel_var_dict["holes_x"][ii], self.kernel_var_dict["holes_y"][ii]) + ): + # Collect sources for this hole and apply weights + + map_of_holes[xx - kx2 : xx + kx2 + adjx, yy - ky2 : yy + ky2 + adjy, :] = ( + hole_idx + ) + S = kspace[xx - kx2 : xx + kx2 + adjx, yy - ky2 : yy + ky2 + adjy, :] + # if patch_index < 10: + # print(f"Kernel-patch {ii} from k-space: {S.shape}") + # imwrite(f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/kernel_kspace_{ii}_hole{xx}_hole{yy}.tiff", S) + # imwrite(f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/abs_kernel_kspace_{ii}_hole{xx}_hole{yy}.tiff", np.abs(S)) + S = S[self.kernel_var_dict["patches"][ii, ...]] + # if patch_index >10: + # print(f"Sources for hole x/y in kspace for patch {ii}: {S.shape}") + # print(f"Weights for hole x/y in kspace for patch {ii}: {weights[ii].shape}") + # imwrite(f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/patch_mask_{ii}.tiff", self.kernel_var_dict['patches'][ii, ...]) + # imwrite(f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/kernel_patch_kspace_{ii}_hole{xx}_hole{yy}.tiff", S) + # imwrite(f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/abs_kernel_patch_kspace_{ii}_hole{xx}_hole{yy}.tiff", np.abs(S)) + # imwrite(f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/weights_{ii}_hole{xx}_hole{yy}.tiff", weights[ii]) + # imwrite(f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/abs_weights_{ii}_hole{xx}_hole{yy}.tiff", np.abs(weights[ii])) + + recon[xx, yy, :] = (weights[ii] @ S[:, None]).squeeze() + # if patch_index > 10: + # print(f"Reconstructed hole x/y in kspace{ii}: {recon[xx, yy, :].shape}") + # imwrite(f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/recon_{ii}_hole{xx}_hole{yy}.tiff", np.moveaxis(recon[xx, yy, :], -1, 0)) + # imwrite(f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/abs_recon_{ii}_hole{xx}_hole{yy}.tiff", np.moveaxis(np.abs(recon[xx, yy, :]), -1,0)) + + imwrite( + f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/map_of_holes_patch_{ii}.tiff", + map_of_holes, + ) + print(f"recon shape before adding to kspace: {recon.shape}") + print(f"(padded) kspace shape before adding recon: {kspace.shape}") + return np.moveaxis((recon[:] + kspace)[kx2:-kx2, ky2:-ky2, :], -1, self.coil_axis) + + +def _kspace_to_log_magnitude(kspace: np.ndarray) -> np.ndarray: + """Convert k-space tensor to a log-magnitude image for visualization.""" + + magnitude = np.log1p(np.abs(kspace)) + + lower = np.quantile(magnitude, 0.05) + upper = np.quantile(magnitude, 0.995) + if float(upper) > float(lower): + magnitude = np.clip(magnitude, lower, upper) + magnitude = (magnitude - lower) / (upper - lower) + else: + mag_max = float(magnitude.max()) + if mag_max > 0.0: + magnitude = magnitude / mag_max + + return np.sqrt(magnitude) + + +if __name__ == "__main__": + filename = "/home/melanie.dohmen/mri_recon/data/fastmri/fastMRI_prostate_T2_IDS_001_020/file_prostate_AXT2_001.h5" + + os.makedirs("/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/", exist_ok=True) + os.makedirs( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/grappa/", exist_ok=True + ) + with h5py.File(filename, "r") as hf: + kspace_data = hf["kspace"][:] + calibration_data = hf["calibration_data"][:] + hdr = hf["ismrmrd_header"][()] + im_recon = hf["reconstruction_rss"][:] + atts = dict() + atts["max"] = hf.attrs["max"] + atts["norm"] = hf.attrs["norm"] + atts["patient_id"] = hf.attrs["patient_id"] + atts["acquisition"] = hf.attrs["acquisition"] + + # (A, S, C, RO, PE) + num_avg, num_slices, num_coils, num_ro, num_pe = kspace_data.shape + + # Calib_data shape: num_slices, num_coils, num_pe_cal + grappa_weight_dict = {} + grappa_weight_dict_2 = {} + + # (A, S, C, RO, PE) -> take first average and slice to get (C, RO, PE) for GRAPPA weight calculation + kspace_slice_regridded = kspace_data[0, 0, ...] + + print("kspace_slice_regridded shape: (A, S, C, RO, PE)") + print(kspace_slice_regridded.shape) + + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_av0_slice0.tiff", + _kspace_to_log_magnitude(kspace_data[0, 0, :, :, :]), + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_av1_slice0.tiff", + _kspace_to_log_magnitude(kspace_data[1, 0, :, :, :]), + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_av2_slice0.tiff", + _kspace_to_log_magnitude(kspace_data[2, 0, :, :, :]), + ) + + print("kspace_slice_regridded transposed shape for GRAPPA: (PE, C, RO)") + print(np.transpose(kspace_slice_regridded, (2, 0, 1)).shape) + + grappa_obj = Grappa( + np.transpose(kspace_slice_regridded, (2, 0, 1)), kernel_size=(5, 5), coil_axis=1 + ) + + kspace_slice_regridded_2 = kspace_data[1, 0, ...] + grappa_obj_2 = Grappa( + np.transpose(kspace_slice_regridded_2, (2, 0, 1)), kernel_size=(5, 5), coil_axis=1 + ) + + # calculate GRAPPA weights + for slice_num in range(num_slices): + # (S, C, PE, cal) -> (C, PE, cal) + calibration_regridded = calibration_data[slice_num, ...] + # (C, PE, cal) -> (cal, C, PE) for GRAPPA weight calculation# + if slice_num == 0: + print(f"calibration_data shape (S, C, PE, cal): {calibration_data.shape}") + print(f"calibration_regridded shape (C, PE, cal): {calibration_regridded.shape}") + print( + f"calibration_regridded transposed shape for GRAPPA: (cal, C, PE)?: {np.transpose(calibration_regridded, (2, 0, 1)).shape}" + ) + grappa_weight_dict[slice_num] = grappa_obj.compute_weights( + np.transpose(calibration_regridded, (2, 0, 1)) + ) + grappa_weight_dict_2[slice_num] = grappa_obj_2.compute_weights( + np.transpose(calibration_regridded, (2, 0, 1)) + ) + + # apply GRAPPA weights + kspace_post_grappa_all = np.zeros(shape=kspace_data.shape, dtype=complex) + + for average, grappa_obj, grappa_weight_dict in zip( + [0, 1, 2], + [grappa_obj, grappa_obj_2, grappa_obj], + [grappa_weight_dict, grappa_weight_dict_2, grappa_weight_dict], + ): + for slice_num in range(num_slices): + # (A, S, C, RO, PE) -> (C, RO, PE) for GRAPPA application + kspace_slice_regridded = kspace_data[average, slice_num, ...] + + # apply weights to transposed k-space slice (PE, C, RO) + kspace_post_grappa = grappa_obj.apply_weights( + np.transpose(kspace_slice_regridded, (2, 0, 1)), grappa_weight_dict[slice_num] + ) + # and move axes back to (C, RO, PE) after GRAPPA application + kspace_post_grappa_all[average, slice_num, ...] = np.moveaxis( + np.moveaxis(kspace_post_grappa, 0, 1), 1, 2 + ) + if average == 0 and slice_num == 0: + print( + f"k-space transposed shape: {np.transpose(kspace_slice_regridded, (2, 0, 1)).shape}" + ) + print(f"k-space post GRAPPA shape: {kspace_post_grappa.shape}") + print( + f"k-space post GRAPPA moved axes shape: {np.moveaxis(np.moveaxis(kspace_post_grappa, 0, 1), 1, 2).shape}" + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_pre_grappa_av0_slice0.tiff", + _kspace_to_log_magnitude(kspace_slice_regridded), + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_post_grappa_av0_slice0.tiff", + _kspace_to_log_magnitude(kspace_post_grappa_all[0, 0, :, :, :]), + ) + + # recon image for each average + im = np.zeros((num_avg, num_slices, num_ro, num_ro)) + for average in range(num_avg): + kspace_grappa = kspace_post_grappa_all[average, ...] + kspace_grappa_padded = zero_pad_kspace_hdr(hdr, kspace_grappa) + im[average] = create_coil_combined_im(kspace_grappa_padded) + imwrite( + f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/im_average_{average}.tiff", + im[average], + ) + + im_3d = np.mean(im, axis=0) + # center crop image to 320 x 320 + img_dict = {} + img_dict["reconstruction_rss"] = center_crop_im(im_3d, [320, 320]) + + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/reconstruction_rss.tiff", + img_dict["reconstruction_rss"], + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/given_reconstruction_rss.tiff", + im_recon, + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/calibratin_data.tiff", + calibration_data[15, 10, :, :], + ) + print("num_avg, num_slices, num_coils, num_ro, num_pe") + print("kspace_data.shape:", kspace_data.shape) + print("kspace_grappa.shape:", kspace_grappa.shape) + print("kspace_grappa_padded.shape:", kspace_grappa_padded.shape) + print("kspace_post_grappa_all.shape:", kspace_post_grappa_all.shape) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_slice_15_coil10.tiff", + _kspace_to_log_magnitude(kspace_data[0, 15, 10, :, :]), + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_grappa_slice_15_coil10.tiff", + _kspace_to_log_magnitude(kspace_grappa[15, 10, :, :]), + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_grappa_padded_slice_15_coil10.tiff", + _kspace_to_log_magnitude(kspace_grappa_padded[15, 10, :, :]), + ) + imwrite( + "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_post_grappa_slice_15_coil10.tiff", + _kspace_to_log_magnitude(kspace_post_grappa_all[0, 15, 10, :, :]), + ) + + for average in range(num_avg): + print("average =", average, ": coil_combined_im(kspace_grappe_padded): ", im[average].shape) + imwrite( + f"/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/im_slice_15_average_{average}.tiff", + im[average][15, :, :], + ) + + print("im_3d.shape:", im_3d.shape) + print("num_slices, num_coils, num_pe_cal") + print("calibration_data.shape:", calibration_data.shape) + print("done") diff --git a/pyproject.toml b/pyproject.toml index 611d63c..c619e37 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ dependencies = [ "deepinv>=0.3.8", "fastmri>=0.3.0", "h5py>=3.16.0", + "mat73>= 0.65", "matplotlib>=3.9.0", "nibabel>=5.3.2", "numpy>=2.4.3", diff --git a/uv.lock b/uv.lock index d25083a..30e3765 100644 --- a/uv.lock +++ b/uv.lock @@ -121,6 +121,73 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, ] +[[package]] +name = "artifactlab" +version = "0.1.0" +source = { virtual = "." } +dependencies = [ + { name = "certifi" }, + { name = "deepinv" }, + { name = "fastmri" }, + { name = "h5py" }, + { name = "mat73" }, + { name = "matplotlib" }, + { name = "nibabel" }, + { name = "numpy" }, + { name = "ptwt" }, + { name = "pydicom" }, + { name = "pytest" }, + { name = "python-certifi-win32" }, + { name = "sigpy" }, + { name = "torch", version = "2.11.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "torchmetrics" }, + { name = "torchvision", version = "0.26.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" }, + { name = "torchvision", version = "0.26.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torchvision", version = "0.26.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "tqdm" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pre-commit" }, + { name = "ruff" }, +] + +[package.metadata] +requires-dist = [ + { name = "certifi", specifier = ">=2026.2.25" }, + { name = "deepinv", specifier = ">=0.3.8" }, + { name = "fastmri", specifier = ">=0.3.0" }, + { name = "h5py", specifier = ">=3.16.0" }, + { name = "mat73", specifier = ">=0.65" }, + { name = "matplotlib", specifier = ">=3.9.0" }, + { name = "nibabel", specifier = ">=5.3.2" }, + { name = "numpy", specifier = ">=2.4.3" }, + { name = "ptwt", specifier = ">=1.0.1" }, + { name = "pydicom", specifier = ">=3.0.1" }, + { name = "pytest", specifier = ">=9.0.2" }, + { name = "python-certifi-win32", specifier = ">=1.6.1" }, + { name = "sigpy", specifier = ">=0.1.27" }, + { name = "torch", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", specifier = ">=2.11.0" }, + { name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.11.0", index = "https://download.pytorch.org/whl/cpu" }, + { name = "torch", marker = "sys_platform == 'linux'", specifier = ">=2.11.0", index = "https://download.pytorch.org/whl/cu128" }, + { name = "torch", marker = "sys_platform == 'win32'", specifier = ">=2.11.0", index = "https://download.pytorch.org/whl/cu128" }, + { name = "torchmetrics", specifier = ">=1.9.0" }, + { name = "torchvision", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", specifier = ">=0.26.0" }, + { name = "torchvision", marker = "sys_platform == 'darwin'", specifier = ">=0.26.0", index = "https://download.pytorch.org/whl/cpu" }, + { name = "torchvision", marker = "sys_platform == 'linux'", specifier = ">=0.26.0", index = "https://download.pytorch.org/whl/cu128" }, + { name = "torchvision", marker = "sys_platform == 'win32'", specifier = ">=0.26.0", index = "https://download.pytorch.org/whl/cu128" }, + { name = "tqdm", specifier = ">=4.67.3" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "pre-commit", specifier = ">=4.2.0" }, + { name = "ruff", specifier = ">=0.11.0" }, +] + [[package]] name = "attrs" version = "25.4.0" @@ -874,6 +941,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" }, ] +[[package]] +name = "mat73" +version = "0.65" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h5py" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/92/61/0e6375513085b13ad23ab150fc83d4d8faa91cba56eb2f30b646259f8214/mat73-0.65.tar.gz", hash = "sha256:ad38a06af3d483632bd939ee572b3724ea8c03d37916765d7278f9de95541ade", size = 19355, upload-time = "2024-07-24T09:12:33.007Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c9/24/e867b1b89b2a2102a5a3bb64ddcd49c5cb815244b2dadff7740c6a422e4f/mat73-0.65-py3-none-any.whl", hash = "sha256:aadfcd00f328eb8f75dd1d4a060a956dc0abefcf5af20f5bc69a5aae64d62cbf", size = 19665, upload-time = "2024-07-24T09:12:31.976Z" }, +] + [[package]] name = "matplotlib" version = "3.10.8" @@ -937,71 +1017,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, ] -[[package]] -name = "mri-recon" -version = "0.1.0" -source = { virtual = "." } -dependencies = [ - { name = "certifi" }, - { name = "deepinv" }, - { name = "fastmri" }, - { name = "h5py" }, - { name = "matplotlib" }, - { name = "nibabel" }, - { name = "numpy" }, - { name = "ptwt" }, - { name = "pydicom" }, - { name = "pytest" }, - { name = "python-certifi-win32" }, - { name = "sigpy" }, - { name = "torch", version = "2.11.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" }, - { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, - { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, - { name = "torchmetrics" }, - { name = "torchvision", version = "0.26.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" }, - { name = "torchvision", version = "0.26.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, - { name = "torchvision", version = "0.26.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, - { name = "tqdm" }, -] - -[package.dev-dependencies] -dev = [ - { name = "pre-commit" }, - { name = "ruff" }, -] - -[package.metadata] -requires-dist = [ - { name = "certifi", specifier = ">=2026.2.25" }, - { name = "deepinv", specifier = ">=0.3.8" }, - { name = "fastmri", specifier = ">=0.3.0" }, - { name = "h5py", specifier = ">=3.16.0" }, - { name = "matplotlib", specifier = ">=3.9.0" }, - { name = "nibabel", specifier = ">=5.3.2" }, - { name = "numpy", specifier = ">=2.4.3" }, - { name = "ptwt", specifier = ">=1.0.1" }, - { name = "pydicom", specifier = ">=3.0.1" }, - { name = "pytest", specifier = ">=9.0.2" }, - { name = "python-certifi-win32", specifier = ">=1.6.1" }, - { name = "sigpy", specifier = ">=0.1.27" }, - { name = "torch", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", specifier = ">=2.11.0" }, - { name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.11.0", index = "https://download.pytorch.org/whl/cpu" }, - { name = "torch", marker = "sys_platform == 'linux'", specifier = ">=2.11.0", index = "https://download.pytorch.org/whl/cu128" }, - { name = "torch", marker = "sys_platform == 'win32'", specifier = ">=2.11.0", index = "https://download.pytorch.org/whl/cu128" }, - { name = "torchmetrics", specifier = ">=1.9.0" }, - { name = "torchvision", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", specifier = ">=0.26.0" }, - { name = "torchvision", marker = "sys_platform == 'darwin'", specifier = ">=0.26.0", index = "https://download.pytorch.org/whl/cpu" }, - { name = "torchvision", marker = "sys_platform == 'linux'", specifier = ">=0.26.0", index = "https://download.pytorch.org/whl/cu128" }, - { name = "torchvision", marker = "sys_platform == 'win32'", specifier = ">=0.26.0", index = "https://download.pytorch.org/whl/cu128" }, - { name = "tqdm", specifier = ">=4.67.3" }, -] - -[package.metadata.requires-dev] -dev = [ - { name = "pre-commit", specifier = ">=4.2.0" }, - { name = "ruff", specifier = ">=0.11.0" }, -] - [[package]] name = "multidict" version = "6.7.1" From dfd99525d170e39b1eb4e07dc9d1fec29504c734 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Mon, 1 Jun 2026 06:03:07 +0000 Subject: [PATCH 02/22] run all with config and all datasets --- examples/config.yaml | 46 ++ examples/config_md.yaml | 46 ++ examples/run_all.py | 755 ++++++++------------------ mri_recon/distortions/__init__.py | 1 + mri_recon/distortions/base.py | 23 + mri_recon/reconstruction/__init__.py | 2 +- mri_recon/reconstruction/inference.py | 50 +- mri_recon/utils/__init__.py | 1 + mri_recon/utils/oasis_adapter.py | 18 +- mri_recon/utils/plot.py | 19 + mri_recon/utils/prostate_adaptor.py | 443 +-------------- 11 files changed, 408 insertions(+), 996 deletions(-) create mode 100644 examples/config.yaml create mode 100644 examples/config_md.yaml diff --git a/examples/config.yaml b/examples/config.yaml new file mode 100644 index 0000000..eaa8040 --- /dev/null +++ b/examples/config.yaml @@ -0,0 +1,46 @@ +data: + "fastmri_knee": "/path/to/fastmri/singlecoil_val" + "oasis": "/path/to/oasis" + "fastmri_brain": "/path/to/fastmri/fastMRI_multicoil_brain_test" + "cmrxrecon": "/path/to/CMRxRecon" + "fastmri_prostate": "/path/to/fastmri/fastMRI_prostate_T2_IDS_001_020" + +distortions: + - "BaseDistortion" + - "CartesianUndersamplingVariableDensity" + - "CartesianUndersamplingUniformRandom" + - "CartesianUndersamplingUniformRandomZeroACS" + - "CartesianUndersamplingEquispaced" + - "CartesianUndersamplingEquispacedZeroACS" + - "PartialFourier" + - "Phase-EncodeGhosting" + - "SegmentedTranslationMotion" + - "SegmentedRotationalMotion" + - "TranslationMotion" + - "RotationalMotion" + - "Off-centerAnisotropicGaussianBiasField" + - "GaussianBiasField" + - "AnisotropicLP" + - "HannTaperLP" + - "KaiserTaperLP" + - "GaussianNoise" + - "IsotropicLP" + - "RadialHigh-passEmphasis" +reconstruction_algorithms: + - "zero-filled" + - "conjugate-gradient" + #- "ram", + #- "dip", + #- "tv-pgd", + #- "wavelet-fista", + #- "tv-fista", + #- "tv-pdhg", + - "unet-fastmri" + - "unet-oasis-acceleration4" + #- "unet-oasis-acceleration8" + #- "unet-oasis-acceleration10" +num_samples: 1 +keep_fraction: 0.25 +center_fraction: 0.125 +verbose: true +results_dir: "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1" \ No newline at end of file diff --git a/examples/config_md.yaml b/examples/config_md.yaml new file mode 100644 index 0000000..ec94872 --- /dev/null +++ b/examples/config_md.yaml @@ -0,0 +1,46 @@ +data: + "fastmri_knee": "/home/melanie.dohmen/ArtifactLab/data/singlecoil_val" + "oasis": "/home/melanie.dohmen/ArtifactLab/data/oasis" + "fastmri_brain": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_multicoil_brain_test" + "cmrxrecon": "/home/melanie.dohmen/ArtifactLab/data/CMRxRecon" + "fastmri_prostate": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_prostate_T2_IDS_001_020" + +distortions: + #- "BaseDistortion" + #- "CartesianUndersamplingVariableDensity" + #- "Cartesian undersampling (uniform random)", + #- "Cartesian undersampling (uniform random, zero ACS)", + #- "Cartesian undersampling (equispaced)", + #- "CartesianUndersamplingEquispacedZeroACS" + #- "Partial Fourier", + #- "Phase-encode ghosting", + #- "Segmented translation motion", + # "Segmented rotational motion", + # "Translation motion", + # "Rotational motion", + # "Off-center anisotropic Gaussian bias field", + # "Gaussian bias field", + # "Anisotropic LP", + # "Hann taper LP", + # "Kaiser taper LP", + - "GaussianNoise" + # "Isotropic LP", + # "Radial high-pass emphasis", +reconstruction_algorithms: + - "zero-filled" + - "conjugate-gradient" + #- "ram", + #- "dip", + #- "tv-pgd", + #- "wavelet-fista", + #- "tv-fista", + #- "tv-pdhg", + - "unet-fastmri" + - "unet-oasis-acceleration4" + #- "unet-oasis-acceleration8" + #- "unet-oasis-acceleration10" +num_samples: 1 +keep_fraction: 0.25 +center_fraction: 0.125 +verbose: true +results_dir: "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1" \ No newline at end of file diff --git a/examples/run_all.py b/examples/run_all.py index 205d171..dbb4f3a 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -10,37 +10,22 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import numpy as np -from pathlib import Path + import deepinv as dinv import torch +import yaml from tifffile import imwrite from mri_recon.distortions import ( - AnisotropicResolutionReduction, BaseDistortion, - CartesianUndersampling, DistortedKspaceMultiCoilMRI, - GaussianKspaceBiasField, - GaussianNoiseDistortion, - HannTaperResolutionReduction, - IsotropicResolutionReduction, - KaiserTaperResolutionReduction, - OffCenterAnisotropicGaussianKspaceBiasField, - PartialFourierDistortion, - PhaseEncodeGhostingDistortion, - RadialHighPassEmphasisDistortion, - RotationalMotionDistortion, - SegmentedRotationalMotionDistortion, - SegmentedTranslationMotionDistortion, - TranslationMotionDistortion, + choose_distortion, ) from mri_recon.reconstruction import ( ConjugateGradientReconstructor, choose_reconstructor, uses_oasis_centered_path, - validate_algorithm_dataset_compatibility, - EXPLICIT_UNET_ALGORITHMS, + compatible_dataset_with_reconstructor, ) from mri_recon.utils import ( OasisCenteredFFTPhysics, @@ -51,279 +36,9 @@ oasis_kspace_to_fastmri_measurement, image_to_kspace, _kspace_to_log_magnitude, + convert_image_for_save, ) -EXPERIMENTS_DIR = Path("reports") / "experiments" - -ALGORITHMS = [ - # "zero-filled", - "conjugate-gradient", - # "ram", - # "dip", - # "tv-pgd", - # "wavelet-fista", - # "tv-fista", - # "tv-pdhg", - *list(EXPLICIT_UNET_ALGORITHMS), -] - -DISTORTIONS = [ - "no distortion", - # "Cartesian undersampling (variable density)", - # "Cartesian undersampling (uniform random)", - # "Cartesian undersampling (uniform random, zero ACS)", - # "Cartesian undersampling (equispaced)", - "Cartesian undersampling (equispaced, zero ACS)", - # "Partial Fourier", - # "Phase-encode ghosting", - # "Segmented translation motion", - # "Segmented rotational motion", - # "Translation motion", - # "Rotational motion", - # "Off-center anisotropic Gaussian bias field", - # "Gaussian bias field", - # "Anisotropic LP", - # "Hann taper LP", - # "Kaiser taper LP", - # "Gaussian noise", - # "Isotropic LP", - # "Radial high-pass emphasis", -] -METRICS = [ - "PSNR", - # "NMSE", - # "SSIM", - # "HaarPSI", - # "SharpnessIndex", - # "BlurStrength", -] - -DATASETS = { - # "fastmri": "/home/melanie.dohmen/mri_recon/data/fastmri/singlecoil_val", - "oasis": "/home/melanie.dohmen/mri_recon/data/oasis", - # "fastmri_multicoil": "/home/melanie.dohmen/mri_recon/data/fastmri/multicoil_train", - "cmrxrecon": "/home/melanie.dohmen/mri_recon/data/CMRxRecon/CMRxRecon/", # SingleCoil/Cine/TrainingSet/FullSample", - "prostate": "/home/melanie.dohmen/mri_recon/data/fastmri/fastMRI_prostate_T2_IDS_001_020", -} - - -def convert_image_for_save(im: torch.Tensor) -> np.ndarray: - """ - Convert a PyTorch tensor image complex tensor to a real-valued NumPy array suitable - by calculating the magnitude. - (B, 2, H, W) or (B, H, W) with complex type -> (B, H, W) - - Args: - im (torch.Tensor): The input image tensor. - - Returns: - np.ndarray: The converted image array. - """ - if torch.is_complex(im) or im.shape[1] == 2: - im = dinv.utils.signals.complex_abs(im, dim=1, keepdim=False) - return im.numpy() - - -def choose_distortion( - name: str, - keep_fraction: float = 0.25, - center_fraction: float = 0.125, - cartesian_axis: int = -2, -) -> BaseDistortion: - """Build one distortion operator for the inference comparison script. - - The ``cartesian_axis`` is supplied by the active measurement convention: - FastMRI-native runs use the repository's existing axis, while OASIS-native - and FastMRI-to-OASIS runs use the centered OASIS axis. - """ - - match name: - case "Phase-encode ghosting": - return PhaseEncodeGhostingDistortion( - line_period=2, - line_offset=1, - phase_error_radians=torch.pi / 2, - corrupted_line_scale=1.0, - ) - case "Cartesian undersampling (variable density)": - return CartesianUndersampling( - keep_fraction=keep_fraction, - center_fraction=center_fraction, - pattern="variable_density_random", - axis=cartesian_axis, - seed=42, - ) - case "Cartesian undersampling (uniform random)": - return CartesianUndersampling( - keep_fraction=keep_fraction, - center_fraction=center_fraction, - pattern="uniform_random", - axis=cartesian_axis, - seed=42, - ) - case "Cartesian undersampling (uniform random, zero ACS)": - return CartesianUndersampling( - keep_fraction=keep_fraction, - center_fraction=0.0, - pattern="uniform_random", - axis=cartesian_axis, - seed=42, - ) - case "Cartesian undersampling (equispaced)": - return CartesianUndersampling( - keep_fraction=keep_fraction, - center_fraction=center_fraction, - pattern="equispaced", - axis=cartesian_axis, - seed=42, - ) - case "Cartesian undersampling (equispaced, zero ACS)": - return CartesianUndersampling( - keep_fraction=keep_fraction, - center_fraction=0.0, - pattern="equispaced", - axis=cartesian_axis, - seed=42, - ) - case "Partial Fourier": - return PartialFourierDistortion( - partial_fraction=0.7, - center_fraction=center_fraction, - axis=cartesian_axis, - side="high", - ) - case "Anisotropic LP": - return AnisotropicResolutionReduction( - kx_radius_fraction=1.0, - ky_radius_fraction=0.25, - ) - case "Hann taper LP": - return HannTaperResolutionReduction( - radius_fraction=0.35, - transition_fraction=0.4, - ) - case "Kaiser taper LP": - return KaiserTaperResolutionReduction( - radius_fraction=0.35, - transition_fraction=0.4, - beta=8.6, - ) - case "Radial high-pass emphasis": - return RadialHighPassEmphasisDistortion(alpha=0.4) - case "Isotropic LP": - return IsotropicResolutionReduction(radius_fraction=0.1) - case "Off-center anisotropic Gaussian bias field": - return OffCenterAnisotropicGaussianKspaceBiasField( - width_x_fraction=0.2, - width_y_fraction=0.35, - center_x_fraction=0.15, - center_y_fraction=-0.1, - edge_gain=0.3, - ) - case "Translation motion": - return TranslationMotionDistortion(shift_x_pixels=60, shift_y_pixels=10) - case "Rotational motion": - return RotationalMotionDistortion(angle_radians=torch.pi / 6) - case "Segmented rotational motion": - return SegmentedRotationalMotionDistortion( - angle_radians=(0.0, torch.pi / 20, -torch.pi / 24, torch.pi / 16), - ) - case "Segmented translation motion": - return SegmentedTranslationMotionDistortion( - shift_x_pixels=(0.0, 20.0, 50.0, -50.0), - shift_y_pixels=(0.0, 10.0, -20.0, 20.0), - ) - case "Gaussian bias field": - return GaussianKspaceBiasField(width_fraction=0.35, edge_gain=0.4) - case "Gaussian noise": - return GaussianNoiseDistortion(sigma=0.00001) - case "no distortion": - return BaseDistortion() - case _: - raise ValueError(f"Unknown distortion {name!r}") - - -def choose_metric(name: str) -> dinv.metric.Metric: - """Build one evaluation metric used in the saved comparison plots.""" - - match name: - case "PSNR": - return dinv.metric.PSNR(max_pixel=None, complex_abs=True) - case "NMSE": - return dinv.metric.NMSE(complex_abs=True) - case "SSIM": - return dinv.metric.SSIM(max_pixel=None, complex_abs=True) - case "HaarPSI": - return dinv.metric.HaarPSI(norm_inputs="min_max", complex_abs=True) - case "BlurStrength": - return dinv.metric.BlurStrength(complex_abs=True) - case "SharpnessIndex": - return dinv.metric.SharpnessIndex(complex_abs=True) - - -# def prepare_measurement_sample( -# sample_batch: object, -# dataset_name: str, -# use_oasis_fft_path: bool, -# run_device: torch.device | str, -# ) -> tuple[torch.Tensor | None, torch.Tensor]: -# """Prepare one input measurement and its clean image reference. - -# FastMRI samples are loaded as native measurements. When the OASIS U-Net is -# selected on FastMRI data, the helper converts those measurements into the -# centered OASIS k-space convention while preserving the native adjoint image -# as the clean reference. -# """ - -# if dataset_name == "oasis": -# x = sample_batch["x"].to(run_device) -# y = image_to_kspace(x) -# coil_maps = None -# elif dataset_name in ("fastmri",) and use_oasis_fft_path: -# y = sample_batch[1].to(run_device) -# x = fastmri_measurement_to_image(y) -# y = fastmri_measurement_to_oasis_kspace(y, device=run_device) -# coil_maps = None -# elif dataset_name == "fastmri_multicoil" and use_oasis_fft_path: -# y = sample_batch[1].to(run_device) - -# coil_maps = ( -# sample_batch[2]["coil_maps"].to(run_device) -# if isinstance(sample_batch, (tuple, list)) -# and len(sample_batch) == 3 -# and "coil_maps" in sample_batch[2] -# else None -# ) -# x = fastmri_measurement_to_image(y, coil_maps=coil_maps) -# y = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) -# elif dataset_name in ("fastmri", "fastmri_multicoil"): -# x = None -# y = sample_batch[1].to(run_device) -# coil_maps = ( -# sample_batch[2]["coil_maps"].to(run_device) -# if isinstance(sample_batch, (tuple, list)) -# and len(sample_batch) == 3 -# and "coil_maps" in sample_batch[2] -# else None -# ) -# elif dataset_name in ("cmrxrecon"): -# x = sample_batch[0].to(run_device) -# y = sample_batch[1].to(run_device) -# coil_maps = ( -# sample_batch[2]["coil_maps"].to(run_device) -# if isinstance(sample_batch, (tuple, list)) -# and len(sample_batch) == 3 -# and "coil_maps" in sample_batch[2] -# else None -# ) -# elif dataset_name in ("prostate"): -# x = sample_batch[0].to(run_device) -# y = sample_batch[1].to(run_device) -# coil_maps = None - -# print(f"\t[Prepared measurement] k-space shape {y.shape} and reference image shape: {x.shape if x is not None else None}") - -# return x, y, coil_maps def get_measurement_sample( @@ -333,24 +48,34 @@ def get_measurement_sample( ) -> tuple[torch.Tensor | None, torch.Tensor]: """Prepare one input measurement and its clean image reference. - FastMRI samples are loaded as native measurements. When the OASIS U-Net is - selected on FastMRI data, the helper converts those measurements into the - centered OASIS k-space convention while preserving the native adjoint image - as the clean reference. + Always prepare a (fast-MRI-like) non-centered k-space measurement + as well as a (oasis-like) centered k-space version of the measurement + and a reference reconstruction in the image domain. """ coil_maps = None if dataset_name == "oasis": + # reference image, shape: (B, 2, H, W) dtype: float32 x = sample_batch["x"].to(run_device) - print(f"\t[Debug] Reference image shape: {x.shape}, dtype: {x.dtype}") + # centered k-space data, shape: (B, 2, H, W) dtype: float32 y_centered = image_to_kspace(x) - print(f"\t[Debug] Centered k-space shape: {y_centered.shape}, dtype: {y_centered.dtype}") + # k-space data, shape: (B, 2, H, W) dtype: float32 y = oasis_kspace_to_fastmri_measurement(y_centered, device=run_device) - elif dataset_name in ("fastmri"): + elif dataset_name == "fastmri_knee": + # reference image, shape: (B, 1, H/2, H/2) dtype: float32 + x = sample_batch[0].to(run_device) + # kspace data, shape: (B, 2, H, W) dtype: float32 y = sample_batch[1].to(run_device) + # centered k-space data, shape: (B, 2, H, W) dtype: float32 y_centered = fastmri_measurement_to_oasis_kspace(y, device=run_device) - x = fastmri_measurement_to_image(y, rss=True) - elif dataset_name in ("fastmri_multicoil"): + # reconstructed reference image: + # shape: (B, 1, H, W) dtype: float32 + print("stop for testing") + elif dataset_name == "fastmri_brain": + # reference image, shape: (B, 1, H/2, H/2) dtype: float32 + x = sample_batch[0].to(run_device) + # kspace data, shape: (B, 2, num_coils, H, W) dtype: float32 y = sample_batch[1].to(run_device) + # coil maps, shape: (B, num_coils, H, W) dtype: complex64 coil_maps = ( sample_batch[2]["coil_maps"].to(run_device) if isinstance(sample_batch, (tuple, list)) @@ -358,15 +83,19 @@ def get_measurement_sample( and "coil_maps" in sample_batch[2] else None ) - x = fastmri_measurement_to_image(y, coil_maps=coil_maps, rss=True) + # centered k-space data, shape: (B, 2, H, W) dtype: float32 y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) - elif dataset_name in ("cmrxrecon"): - # ignore multi-coil reference image data + + elif dataset_name == "cmrxrecon": + # reference image, shape: (B, 2, n_timepoints, (n_coils), H, W) x = sample_batch[0].to(run_device) print(f"\t[Debug] Reference image shape: {x.shape}, dtype: {x.dtype}") - + # k-space data, shape: (B, 2, n_timepoints, (n_coils), H, W) dtype: float32 y = sample_batch[1].to(run_device) + print(f"\t[Debug] k-space shape: {y.shape}, dtype: {y.dtype}") + # not available for all samples, either None or + # shape (1, num_coils, H, W) coil_maps = ( sample_batch[2]["coil_maps"].to(run_device) if isinstance(sample_batch, (tuple, list)) @@ -374,112 +103,66 @@ def get_measurement_sample( and "coil_maps" in sample_batch[2] else None ) + print(f"\t[Debug] Coil maps shape: {coil_maps.shape if coil_maps is not None else None}, dtype: {coil_maps.dtype if coil_maps is not None else None}") + # centered k-space data, shape: (B, 2, H, W) dtype: float32 y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) + print(f"\t[Debug] Centered k-space shape: {y_centered.shape}, dtype: {y_centered.dtype}") # reconstruct coil-combined image reference from multi-coil k-space data using # integrated espirit sensitivity map estimation, RSS coil combination - x = fastmri_measurement_to_image(y, coil_maps=coil_maps, rss=True) - - elif dataset_name in ("prostate"): - # has shape: (B, num_averages, coils, H, W) with dtype= complex128-> take first average and convert to image space reference + + #x = fastmri_measurement_to_image(y, coil_maps=coil_maps, rss=True) + #print(f"\t[Debug] Reference image shape: {x.shape}, dtype: {x.dtype}") + elif dataset_name == "fastmri_prostate": + # reference image, shape: (slices, W, H): dtype float32 x = sample_batch[0].to(run_device) print(f"\t[Debug] Reference image shape: {x.shape}, type: {x.dtype}") - # take mean of average images: - x = x.mean(dim=1) - print(f"\t[Debug] Mean of averages image shape: {x.shape}, type: {x.dtype}") - - # convert to channel representation of complex numbers - # (B, H, W) with complex dtype -> (B, H, W, 2) with real dtype - x = torch.view_as_real(x) if torch.is_complex(x) else x - print(f"\t[Debug] after view_as_real (if complex) shape: {x.shape}, type: {x.dtype}") - # move channel with real and imaginary parts to channel dimension - # (B, H, W, 2) -> (B, 2, H, W) - x = x.moveaxis(-1, 1) - print(f"\t[Debug] after moving channels: {x.shape}, type: {x.dtype}") + # add zero imaginary channel: + # (B, slices, H, W) -> (B, 2, slices, H, W) + x = torch.stack([x, torch.zeros_like(x)], dim=1) + - # ignore k-space data: - y = sample_batch[1].to(run_device) - print(f"\t[Debug] Original k-space shape: {y.shape}, type: {y.dtype}") + # y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) # create oasis-like k-space data from image: y_centered = image_to_kspace(x) print(f"\t[Debug] Centered k-space shape: {y_centered.shape}, type: {y_centered.dtype}") y = oasis_kspace_to_fastmri_measurement(y_centered, device=run_device) - print(f"\tk-space shape {y.shape} and reference image shape: {x.shape}") + print(f"\tk-space shape {y.shape}[{y.dtype}] and reference image shape: {x.shape}[{x.dtype}]") + if coil_maps is not None: + print(f"\tcoil maps shape: {coil_maps.shape}[{coil_maps.dtype}]") return x, y, y_centered, coil_maps if __name__ == "__main__": - # parser = argparse.ArgumentParser(description=__doc__) - - # # data related arguments - # parser.add_argument( - # "--source", - # type=Path, - # help="Local FastMRI directory with raw k-space .h5 files or OASIS root directory.", - # ) - # parser.add_argument( - # "--dataset", - # choices=("fastmri", "oasis", "fastmri_multicoil", "cmrxrecon"), - # default="fastmri", - # ) - - # parser.add_argument("--distortion", type=str, default="", choices=DISTORTIONS) - # parser.add_argument( - # "--keep_fraction", - # type=float, - # default=0.25, - # help="Fraction of k-space lines to keep for undersampling distortions.", - # ) - # parser.add_argument( - # "--center_fraction", - # type=float, - # default=0.125, - # help="Fraction of low-frequency k-space lines to keep fully for undersampling distortions.", - # ) - - # # algo related arguments - # parser.add_argument( - # "--algorithm", - # type=str, - # default="", - # choices=ALGORITHMS, - # help="Reconstruction algorithm applied to undistorted and distorted k-space.", - # ) - # # inference related arguments - # parser.add_argument("--num_samples", type=int, default=1, help="How many samples to process.") - # parser.add_argument( - # "--verbose", - # action="store_true", - # help="Enable verbose output for reconstructors that support it.", - # ) - # args = parser.parse_args() - num_samples = 1 - keep_fraction = 0.25 - center_fraction = 0.125 - verbose = True - - os.makedirs(EXPERIMENTS_DIR, exist_ok=True) - - # set up device, dataset, metrics + + # read config file in yaml format as first argument from commmand line + if len(sys.argv) < 2: + print("Usage: python examples/run_all.py ") + sys.exit(1) + + with open(sys.argv[1], "r") as f: + config = yaml.safe_load(f) + + os.makedirs(config["results_dir"], exist_ok=True) + + # set up device device = dinv.utils.get_device() - for dataset_name, dataset_rootdir in DATASETS.items(): + for dataset_name, dataset_rootdir in config["data"].items(): + print(f"=== {dataset_name} ===") - selected_algorithms = ALGORITHMS - selected_distortions = DISTORTIONS - + # initialize dataset if dataset_name == "oasis": - # split_csv = OASISSinglecoilUnetReconstructor.resolve_default_split_csv() dataset = OasisCenterSliceFolderDataset( data_path=dataset_rootdir, ) - elif dataset_name == "fastmri": + elif dataset_name == "fastmri_knee": dataset = dinv.datasets.FastMRISliceDataset(str(dataset_rootdir), slice_index="middle") - elif dataset_name == "fastmri_multicoil": + elif dataset_name == "fastmri_brain": dataset = dinv.datasets.FastMRISliceDataset( str(dataset_rootdir), slice_index="middle", @@ -494,15 +177,16 @@ def get_measurement_sample( data_dir="SingleCoil/Cine/TrainingSet/FullSample", apply_mask=False, ) - elif dataset_name == "prostate": - dataset = FastMRIProstateDataset(data_path=dataset_rootdir, num_samples=num_samples) + elif dataset_name == "fastmri_prostate": + dataset = FastMRIProstateDataset(data_path=dataset_rootdir, num_samples=config["num_samples"], slice_index="middle") else: raise NotImplementedError(f"Invalid dataset: {dataset_name}") - metrics = [choose_metric(m) for m in METRICS] + # loop through samples of dataset for i, batch in enumerate(iter(torch.utils.data.DataLoader(dataset))): + # exit loop if we have processed the specified number of samples - if i >= num_samples: + if i >= config["num_samples"]: break print(f"{dataset_name} sample {i}...") @@ -512,173 +196,188 @@ def get_measurement_sample( run_device=device, ) - physics_clean_oasis_fft_path = OasisCenteredFFTPhysics(BaseDistortion()) - physics_clean_fastmri_path = DistortedKspaceMultiCoilMRI( - BaseDistortion(), img_size=x_reference.shape, coil_maps=coil_maps, device=device - ) - - x_clean_oasis_fft_path = ConjugateGradientReconstructor()( - y, physics_clean_oasis_fft_path - ) - x_clean_fastmri_path = ConjugateGradientReconstructor()(y, physics_clean_fastmri_path) - - # reference reconstructions: + # save reference image imwrite( - os.path.join( - EXPERIMENTS_DIR, f"image_{dataset_name}_sample_{i}_CG_oasis_fft_path.tiff" - ), - convert_image_for_save(x_clean_oasis_fft_path), - ) - imwrite( - os.path.join( - EXPERIMENTS_DIR, f"image_{dataset_name}_sample_{i}_CG_fastmri_path.tiff" - ), - convert_image_for_save(x_clean_fastmri_path), - ) - imwrite( - os.path.join(EXPERIMENTS_DIR, f"image_{dataset_name}_sample_{i}_reference.tiff"), + os.path.join(config["results_dir"], f"image_{dataset_name}_sample_{i}_reference.tiff"), convert_image_for_save(x_reference), ) - for distortion_name in selected_distortions: + # use fast-mri type samples first, later proceed with oasis-centered fft path + physics_clean = DistortedKspaceMultiCoilMRI( + BaseDistortion(), img_size=y.shape[-2:], coil_maps=coil_maps, device=device + ) + + # reference from dataset: + for distortion_name in config["distortions"]: print(f"\t{distortion_name} ...") - distortion_oasis_fft_path = choose_distortion( - distortion_name, - keep_fraction=keep_fraction, - center_fraction=center_fraction, - cartesian_axis=-1, - ) - distortion_fastmri_path = choose_distortion( + distortion = choose_distortion( distortion_name, - keep_fraction=keep_fraction, - center_fraction=center_fraction, + keep_fraction=config["keep_fraction"], + center_fraction=config["center_fraction"], cartesian_axis=-2, ) - y_distorted_oasis_fft_path = distortion_oasis_fft_path.A(y_centered) - y_distorted_fastmri_path = distortion_fastmri_path.A(y) + y_distorted = distortion.A(y) - physics_distorted_oasis_fft_path = OasisCenteredFFTPhysics( - distortion_oasis_fft_path - ) - physics_distorted_fastmri_path = DistortedKspaceMultiCoilMRI( - distortion_fastmri_path, - img_size=x_reference.shape, + physics_distorted = DistortedKspaceMultiCoilMRI( + distortion, + img_size=y.shape[-2:], coil_maps=coil_maps, device=device, ) - for algo_name in selected_algorithms: - print(f"\t\t{algo_name} ...") - try: - validate_algorithm_dataset_compatibility(dataset_name, algo_name) - - use_oasis_path = uses_oasis_centered_path(dataset_name, algo_name) - if use_oasis_path: - y_distorted = y_distorted_oasis_fft_path - physics_clean = physics_clean_oasis_fft_path - physics_distorted = physics_distorted_oasis_fft_path - x_clean = x_reference - else: - y_distorted = y_distorted_fastmri_path - physics_clean = physics_clean_fastmri_path - physics_distorted = physics_distorted_fastmri_path - - algo = choose_reconstructor( - algo_name, - img_size=y_distorted.shape[-2:], - device=device, - verbose=verbose, - dataset=dataset_name, - ).to(device) - - # save reference and distorted k-space for debugging and visualization purposes - imwrite( - os.path.join( - EXPERIMENTS_DIR, f"kspace_{dataset_name}_sample_{i}_reference.tiff" - ), - _kspace_to_log_magnitude(y).numpy(), - ) - imwrite( - os.path.join( - EXPERIMENTS_DIR, - f"kspace_{dataset_name}_sample_{i}_{distortion_name}.tiff", - ), - _kspace_to_log_magnitude(y_distorted).numpy(), - ) - - # actual reconstruction with the algo being evaluated - try: - if dataset_name == "prostate": - # prostate dataset has multiple k-space averages, - # so we reconstruct each average separately and then average in the image domain - x_corrected_averages = [] - x_uncorrected_averages = [] - for average in range(y_distorted.shape[0]): - x_uncorrected_averages.append( - algo(y_distorted[average], physics_clean) - ) - x_corrected_averages.append( - algo(y_distorted[average], physics_distorted) - ) - - x_uncorrected = torch.stack(x_uncorrected_averages, dim=0).mean( - dim=0 - ) - x_corrected = torch.stack(x_corrected_averages, dim=0).mean(dim=0) + for reconstructor_name in config["reconstruction_algorithms"]: + print(f"\t\t{reconstructor_name} ...") + if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): + + # only run on reconstructors, that use the fastmri-like k-space + if not uses_oasis_centered_path(dataset_name, reconstructor_name): - else: - x_uncorrected = algo(y_distorted, physics_clean) - x_corrected = algo(y_distorted, physics_distorted) + reconstructor = choose_reconstructor( + reconstructor_name, + img_size=y_distorted.shape[-2:], + device=device, + verbose=config["verbose"], + ).to(device) - # performed reconstruction images + # save reference and distorted k-space for debugging purposes imwrite( os.path.join( - EXPERIMENTS_DIR, - f"image_{dataset_name}_sample_{i}_{distortion_name}_{algo_name}_uncorrected.tiff", + config["results_dir"], f"kspace_{dataset_name}_sample_{i}_reference.tiff" ), - convert_image_for_save(x_uncorrected), + _kspace_to_log_magnitude(y).numpy(), ) imwrite( os.path.join( - EXPERIMENTS_DIR, - f"image_{dataset_name}_sample_{i}_{distortion_name}_{algo_name}_corrected.tiff", + config["results_dir"], + f"kspace_{dataset_name}_sample_{i}_{distortion_name}.tiff", ), - convert_image_for_save(x_corrected), + _kspace_to_log_magnitude(y_distorted).numpy(), ) - except Exception as e: - print( - f"Error reconstructing algo {algo_name} with distortion {distortion_name} on sample {i}: {e}" + # actual reconstruction with the selected reconstructor + try: + + x_uncorrected = reconstructor(y_distorted, physics_clean) + x_corrected = reconstructor(y_distorted, physics_distorted) + + + # crop recostructed image to reference image size: + if x_uncorrected.shape[-2:] != x_reference.shape[-2:]: + x_uncorrected = physics_clean.crop(x_uncorrected, shape=x_reference.shape[-2:]) + + if x_corrected.shape[-2:] != x_reference.shape[-2:]: + x_corrected_clean = physics_distorted.crop(x_corrected, shape=x_reference.shape[-2:]) + + # save reconstructed images + imwrite( + os.path.join( + config["results_dir"], + f"image_{dataset_name}_sample_{i}_{distortion_name}_{reconstructor_name}_uncorrected.tiff", + ), + convert_image_for_save(x_uncorrected), + ) + imwrite( + os.path.join( + config["results_dir"], + f"image_{dataset_name}_sample_{i}_{distortion_name}_{reconstructor_name}_corrected.tiff", + ), + convert_image_for_save(x_corrected), + ) + + except Exception as e: + print( + f"Error using {reconstructor_name}: {e}" + ) + + + else: + print(f"\t\t ... not compatible with {dataset_name}") + + + # now proceed with oasis-centered fft path + physics_clean = OasisCenteredFFTPhysics(BaseDistortion()) + + for distortion_name in config["distortions"]: + print(f"\t{distortion_name} ...") + distortion = choose_distortion( + distortion_name, + keep_fraction=config["keep_fraction"], + center_fraction=config["center_fraction"], + cartesian_axis=-1, + ) + + + y_distorted = torch.fft.fftshift(distortion.A(torch.fft.fftshift(y_centered, dim=(-1, -2))), dim=(-2, -1)) + + physics_distorted = OasisCenteredFFTPhysics( + distortion + ) + + for reconstructor_name in config["reconstruction_algorithms"]: + print(f"\t\t{reconstructor_name} ...") + if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): + + # skip all reconstructors, that don't use the oasis-centered path + if uses_oasis_centered_path(dataset_name, reconstructor_name): + + reconstructor = choose_reconstructor( + reconstructor_name, + img_size=y_distorted.shape[-2:], + device=device, + verbose=config["verbose"], + ).to(device) + + # save reference and distorted k-space for debugging purposes + imwrite( + os.path.join( + config["results_dir"], f"kspace_centered_{dataset_name}_sample_{i}_reference.tiff" + ), + _kspace_to_log_magnitude(y_centered).numpy(), + ) + imwrite( + os.path.join( + config["results_dir"], + f"kspace_centered_{dataset_name}_sample_{i}_{distortion_name}.tiff", + ), + _kspace_to_log_magnitude(y_distorted).numpy(), ) - # dinv.utils.plot( - # { - # "Undistorted ksp, CG recon": x_clean, - # "Distorted ksp, CG recon": x_distorted, - # f"Distorted ksp, {algo_name} recon, uncorrected": x_uncorrected, - # f"Distorted ksp, {algo_name} recon, corrected": x_corrected, - # }, - # subtitles=[ - # "", - # "", - # "\n".join( - # f"{m.__class__.__name__} {m(x_uncorrected, x_clean).item():.2f}" - # for m in metrics - # ), - # "\n".join( - # f"{m.__class__.__name__} {m(x_corrected, x_clean).item():.2f}" - # for m in metrics - # ), - # ], - # show=False, - # close=True, - # suptitle=f"Algo {algo_name}, distortion {distortion_name}, Sample {i}", - # save_fn=REPORT_DIR / f"ALGO_{algo_name}_{distortion_name}_sample_{i}.png", - # fontsize=3, - # ) - except Exception as e: - print( - f"\t\tError processing algo {algo_name}, distortion {distortion_name}, sample {i}: {e}" - ) + # actual reconstruction with the algo being evaluated + try: + + x_uncorrected = reconstructor(y_distorted, physics_clean) + x_corrected = reconstructor(y_distorted, physics_distorted) + + if x_uncorrected.shape[-2:] != x_reference.shape[-2:]: + x_uncorrected = physics_clean.crop(x_uncorrected, shape=x_reference.shape[-2:]) + + if x_corrected.shape[-2:] != x_reference.shape[-2:]: + x_corrected_clean = physics_distorted.crop(x_corrected, shape=x_reference.shape[-2:]) + + + # save reconstructed images + imwrite( + os.path.join( + config["results_dir"], + f"image_{dataset_name}_sample_{i}_{distortion_name}_{reconstructor_name}_uncorrected.tiff", + ), + convert_image_for_save(x_uncorrected), + ) + imwrite( + os.path.join( + config["results_dir"], + f"image_{dataset_name}_sample_{i}_{distortion_name}_{reconstructor_name}_corrected.tiff", + ), + convert_image_for_save(x_corrected), + ) + + except Exception as e: + print( + f"\t\tError using {reconstructor_name} with distortion {distortion_name} on sample {i}: {e}" + ) + + else: + print(f"\t\t ... not compatible with {dataset_name}") + diff --git a/mri_recon/distortions/__init__.py b/mri_recon/distortions/__init__.py index af22fd1..1b99916 100644 --- a/mri_recon/distortions/__init__.py +++ b/mri_recon/distortions/__init__.py @@ -20,3 +20,4 @@ RadialHighPassEmphasisDistortion, ) from .undersampling import CartesianUndersampling, PartialFourierDistortion +from .utils import choose_distortion \ No newline at end of file diff --git a/mri_recon/distortions/base.py b/mri_recon/distortions/base.py index 61bc7f3..51d9e0a 100644 --- a/mri_recon/distortions/base.py +++ b/mri_recon/distortions/base.py @@ -180,8 +180,31 @@ def A(self, x: torch.Tensor) -> torch.Tensor: return self.distortion(y) def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: + if len(y.shape) == (5 if self.three_d else 4): y = y.unsqueeze(2) # add coil dim if singlecoil y = self.distortion.A_adjoint(y) + + # # in order to match the reference image shape + # # a crop must be performed here on k-space data: + # if y.shape[-2:] != self.img_size[-2:]: + # y =self.crop(y, crop=True) + + # # and on coil maps for multi-coil data: + # if self.coil_maps is not None and self.coil_maps.shape[2:] != self.img_size[1:]: + # self.coil_maps = self.crop(self.coil_maps, crop=True) return super().A_adjoint(y) + + # def A_dagger(self, y: torch.Tensor, **kwargs) -> torch.Tensor: + + # # in order to match the reference image shape + # # a crop must be performed here on k-space data: + # if y.shape[-2:] != self.img_size[-2:]: + # y =self.crop(y, crop=True) + + # # and on coil maps for multi-coil data: + # if self.coil_maps is not None and self.coil_maps.shape[2:] != self.img_size[1:]: + # self.coil_maps = self.crop(self.coil_maps, crop=True) + + # return super().A_dagger(y, coil_maps=self.coil_maps, **kwargs) diff --git a/mri_recon/reconstruction/__init__.py b/mri_recon/reconstruction/__init__.py index 5eb75ea..fef846c 100644 --- a/mri_recon/reconstruction/__init__.py +++ b/mri_recon/reconstruction/__init__.py @@ -10,7 +10,7 @@ OASIS_UNET_ALGORITHMS, choose_reconstructor, uses_oasis_centered_path, - validate_algorithm_dataset_compatibility, + compatible_dataset_with_reconstructor, ) from .classic import ( ZeroFilledReconstructor, diff --git a/mri_recon/reconstruction/inference.py b/mri_recon/reconstruction/inference.py index 27c1c65..6b9216f 100644 --- a/mri_recon/reconstruction/inference.py +++ b/mri_recon/reconstruction/inference.py @@ -43,31 +43,26 @@ def uses_oasis_centered_path( return algorithm in OASIS_UNET_ALGORITHMS -def validate_algorithm_dataset_compatibility(dataset: str, algorithm: str) -> None: - """Raise a clear error when an explicit algorithm is incompatible with a dataset.""" - - if dataset == "oasis" and algorithm == FASTMRI_UNET_ALGORITHM: - raise ValueError( - "The algorithm 'unet-fastmri' is not supported on the OASIS dataset. " - "Use one of the explicit OASIS U-Net algorithms instead." - ) - elif dataset == "fastmri" and algorithm in OASIS_UNET_ALGORITHMS: - raise ValueError( - "The algorithm 'unet-oasis' is not supported on the FastMRI dataset. " - "Use the 'unet-fastmri' algorithm instead." - ) - elif dataset == "fastmri-multicoil" and algorithm == FASTMRI_UNET_ALGORITHM: - raise ValueError( - "The algorithm 'unet-fastmri' (knee) is not supported on the FastMRI multicoil (brain) dataset. " - "Use the 'unet-oasis' algorithm instead." - ) - elif dataset in ["cmrxrecon", "prostate"] and algorithm in [FASTMRI_UNET_ALGORITHM] + list( - OASIS_UNET_ALGORITHMS.keys() - ): - raise ValueError( - f"The algorithm {algorithm} ({'heart' if dataset == 'cmrxrecon' else 'prostate'}) is not supported on the cmrxrecon or prostate datasets. " - "No trained unet model available for this dataset." - ) +def compatible_dataset_with_reconstructor(dataset: str, reconstructor_name: str) -> bool: + """Check if dataset and trained reconstructor are compatible""" + + # fast mri u-net is only trained with knee data + if reconstructor_name == FASTMRI_UNET_ALGORITHM: + if (dataset == "fastmri_knee"): + return True + else: + return False + + # oasis is only trained with brain data + elif reconstructor_name in OASIS_UNET_ALGORITHMS: + if dataset in ["fastmri_brain", "oasis"]: + return True + else: + return False + + # all other (classic) reconstructors work with any dataset: + else: + return True def choose_reconstructor( @@ -75,7 +70,7 @@ def choose_reconstructor( img_size: tuple = (640, 368), device: torch.device | str = "cpu", verbose: bool = False, - dataset: str = "fastmri", + dataset: str | None = None, ) -> dinv.models.Reconstructor: """Create a reconstructor while enforcing the supported dataset/model matrix. @@ -95,7 +90,8 @@ def choose_reconstructor( explicit algorithm names that are dataset-specific. """ - validate_algorithm_dataset_compatibility(dataset, name) + if dataset is not None and not compatible_dataset_with_reconstructor(dataset, name): + raise ValueError(f"Reconstructor {name} is not compatible with dataset {dataset}, because it was trained with a different image domain.") match name: case "zero-filled": diff --git a/mri_recon/utils/__init__.py b/mri_recon/utils/__init__.py index de656c6..ed72870 100644 --- a/mri_recon/utils/__init__.py +++ b/mri_recon/utils/__init__.py @@ -17,6 +17,7 @@ from .prostate_adaptor import FastMRIProstateDataset as FastMRIProstateDataset from .plot import save_kspace_plot as save_kspace_plot from .plot import _kspace_to_log_magnitude as _kspace_to_log_magnitude +from .plot import convert_image_for_save as convert_image_for_save __all__ = [ "download_file_with_sha256", diff --git a/mri_recon/utils/oasis_adapter.py b/mri_recon/utils/oasis_adapter.py index e87cd8f..89f3ce6 100644 --- a/mri_recon/utils/oasis_adapter.py +++ b/mri_recon/utils/oasis_adapter.py @@ -296,7 +296,10 @@ def fastmri_measurement_to_oasis_kspace( Centered OASIS-convention k-space tensor with shape ``(B, 2, H, W)``. """ - return image_to_kspace(fastmri_measurement_to_image(y, coil_maps=coil_maps, device=device)) + result = image_to_kspace(fastmri_measurement_to_image(y, coil_maps=coil_maps, device=device)) + + + return result def image_to_fast_mri_measurement( @@ -308,8 +311,8 @@ def image_to_fast_mri_measurement( Parameters ---------- - y : torch.Tensor - OASIS k-space measurement tensor with shape ``(B, 2, H, W)``. + x : torch.Tensor + image tensor with shape ``(B, 2, H, W)``. coil_maps : torch.Tensor | None, optional Coil sensitivity maps with shape ``(B, C, H, W)``, where ``C`` is the number of coils. device : torch.device | str, optional @@ -334,8 +337,6 @@ def image_to_fast_mri_measurement( def oasis_kspace_to_fastmri_measurement( y: torch.Tensor, - coil_maps: torch.Tensor | None = None, - device: torch.device | str | None = None, ) -> torch.Tensor: """Adapt OASIS-convention k-space to FastMRI measurement convention. @@ -343,10 +344,7 @@ def oasis_kspace_to_fastmri_measurement( ---------- y : torch.Tensor Centered OASIS-convention k-space tensor with shape ``(B, 2, H, W)``. - coil_maps : torch.Tensor | None, optional - Coil sensitivity maps with shape ``(B, C, H, W)``, where ``C`` is the number of coils. - device : torch.device | str, optional - Device on which to instantiate the temporary native physics operator. + Returns ------- @@ -354,7 +352,7 @@ def oasis_kspace_to_fastmri_measurement( FastMRI-convention k-space tensor with shape ``(B, 2, H, W)``. """ - return image_to_fast_mri_measurement(kspace_to_image(y), coil_maps=coil_maps, device=device) + return image_to_fast_mri_measurement(kspace_to_image(y)) class OasisCenteredFFTPhysics(dinv.physics.LinearPhysics): diff --git a/mri_recon/utils/plot.py b/mri_recon/utils/plot.py index b3ff225..5379519 100644 --- a/mri_recon/utils/plot.py +++ b/mri_recon/utils/plot.py @@ -6,6 +6,8 @@ import matplotlib.pyplot as plt import torch +import numpy as np +import deepinv as dinv def _kspace_to_log_magnitude(kspace: torch.Tensor) -> torch.Tensor: @@ -71,3 +73,20 @@ def save_kspace_plot( ax.axis("off") fig.savefig(save_fn, dpi=200, bbox_inches="tight") plt.close(fig) + + +def convert_image_for_save(im: torch.Tensor) -> np.ndarray: + """ + Convert a PyTorch tensor image complex tensor to a real-valued NumPy array + by calculating the magnitude. + (B, 2, H, W) or (B, H, W) with complex type -> (B, H, W) + + Args: + im (torch.Tensor): The input image tensor. + + Returns: + np.ndarray: The converted image array. + """ + if torch.is_complex(im) or im.shape[1] == 2: + im = dinv.utils.signals.complex_abs(im, dim=1, keepdim=False) + return im.numpy() diff --git a/mri_recon/utils/prostate_adaptor.py b/mri_recon/utils/prostate_adaptor.py index e12f5b8..5145fcd 100644 --- a/mri_recon/utils/prostate_adaptor.py +++ b/mri_recon/utils/prostate_adaptor.py @@ -1,461 +1,44 @@ import os import glob -from typing import Dict, Tuple, Sequence import h5py import numpy as np -from skimage.util import view_as_windows -from tempfile import NamedTemporaryFile as NTF -from tifffile import imwrite import torch -import xml.etree.ElementTree as etree - - -class Grappa: - def __init__( - self, kspace: np.ndarray, kernel_size: Tuple[int, int] = (5, 5), coil_axis: int = -1 - ) -> None: - self.kspace = kspace - self.kernel_size = kernel_size - self.coil_axis = coil_axis - self.lamda = 0.01 - - self.kernel_var_dict = self.get_kernel_geometries() - - def get_kernel_geometries(self): - """ - Extract unique kernel geometries based on a slice of kspace data - - Returns - ------- - geometries : dict - A dictionary containing the following keys: - - 'patches': an array of overlapping patches from the k-space data. - - 'patch_indices': an array of unique patch indices. - - 'holes_x': a dictionary of x-coordinates for holes in each patch. - - 'holes_y': a dictionary of y-coordinates for holes in each patch. - - Notes - ----- - This function extracts unique kernel geometries from a slice of k-space data. - The geometries correspond to overlapping patches that contain at least one hole. - A hole is defined as a region of k-space data where the absolute value of the - complex signal is equal to zero. The function returns a dictionary containing - information about the patches and holes, which can be used to compute weights - for each geometry using the GRAPPA algorithm. - - """ - self.kspace = np.moveaxis(self.kspace, self.coil_axis, -1) - - # Quit early if there are no holes - if np.sum((np.abs(self.kspace[..., 0]) == 0).flatten()) == 0: - return np.moveaxis(self.kspace, -1, self.coil_axis) - - kx, ky = self.kernel_size[:] - kx2, ky2 = int(kx / 2), int(ky / 2) - nc = self.kspace.shape[-1] - - self.kspace = np.pad(self.kspace, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") - - mask = np.ascontiguousarray(np.abs(self.kspace[..., 0]) > 0) - - with NTF() as fP: - # Get all overlapping patches from the mask - P = np.memmap( - fP, - dtype=mask.dtype, - mode="w+", - shape=(mask.shape[0] - 2 * kx2, mask.shape[1] - 2 * ky2, 1, kx, ky), - ) - P = view_as_windows(mask, (kx, ky)) - Psh = P.shape[:] # save shape for unflattening indices later - P = P.reshape((-1, kx, ky)) - - # Find the unique patches and associate them with indices - P, iidx = np.unique(P, return_inverse=True, axis=0) - - # Filter out geometries that don't have a hole at the center. - # These are all the kernel geometries we actually need to - # compute weights for. - validP = np.argwhere(~P[:, kx2, ky2]).squeeze() - - # ignore empty patches - invalidP = np.argwhere(np.all(P == 0, axis=(1, 2))) - validP = np.setdiff1d(validP, invalidP, assume_unique=True) - - validP = np.atleast_1d(validP) - - # Give P back its coil dimension - P = np.tile(P[..., None], (1, 1, 1, nc)) - - holes_x = {} - holes_y = {} - for ii in validP: - # x, y define where top left corner is, so move to ctr, - # also make sure they are iterable by enforcing atleast_1d - idx = np.unravel_index(np.argwhere(iidx == ii), Psh[:2]) - x, y = idx[0] + kx2, idx[1] + ky2 - x = np.atleast_1d(x.squeeze()) - y = np.atleast_1d(y.squeeze()) - - holes_x[ii] = x - holes_y[ii] = y - - return {"patches": P, "patch_indices": validP, "holes_x": holes_x, "holes_y": holes_y} - - def compute_weights(self, calib: np.ndarray) -> Dict[int, np.ndarray]: - """ - Compute the GRAPPA weights for each slice in the input calibration data. - - Parameters: - ---------- - calib : numpy.ndarray - Calibration data with shape (Nx, Nc, Ny) where Nx, Ny are the size of the image in the x and y dimensions, - respectively, and Nc is the number of coils. - - Returns: - ------- - weights : dict - A dictionary of GRAPPA weights for each patch index. - - Notes: - ----- - The GRAPPA algorithm is used to estimate the missing k-space data in undersampled MRI acquisitions. - The algorithm used to compute the GRAPPA weights involves first extracting patches from the calibration data, - and then solving a linear system to estimate the weights. The resulting weights are stored in a dictionary - where the key is the patch index. The equation to solve for the weights involves taking the product of the - sources and the targets in the patch domain, and then regularizing the matrix using Tikhonov regularization. - The function uses numpy's `memmap` to store temporary files to avoid overwhelming memory usage. - """ - - calib = np.moveaxis(calib, self.coil_axis, -1) - kx, ky = self.kernel_size[:] - kx2, ky2 = int(kx / 2), int(ky / 2) - nc = calib.shape[-1] - - calib = np.pad(calib, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") - - # Store windows in temporary files so we don't overwhelm memory - with NTF() as fA: - # Get all overlapping patches of ACS - try: - A = np.memmap( - fA, - dtype=calib.dtype, - mode="w+", - shape=(calib.shape[0] - 2 * kx, calib.shape[1] - 2 * ky, 1, kx, ky, nc), - ) - A[:] = view_as_windows(calib, (kx, ky, nc)).reshape((-1, kx, ky, nc)) - except ValueError: - A = view_as_windows(calib, (kx, ky, nc)).reshape((-1, kx, ky, nc)) - - weights = {} - - for ii in self.kernel_var_dict["patch_indices"]: - # Get the sources by masking all patches of the ACS and - # get targets by taking the center of each patch. Source - # and targets will have the following sizes: - # S : (# samples, N possible patches in ACS) - # T : (# coils, N possible patches in ACS) - # Solve the equation for the weights: using numpy.linalg.solve, - # and Tikhonov regularization for better conditioning: - # SW = T - # S^HSW = S^HT - # W = (S^HS)^-1 S^HT - # -> W = (S^HS + lamda I)^-1 S^HT - - S = A[:, self.kernel_var_dict["patches"][ii, ...]] - T = A[:, kx2, ky2, :] - ShS = S.conj().T @ S - ShT = S.conj().T @ T - lamda0 = self.lamda * np.linalg.norm(ShS) / ShS.shape[0] - weights[ii] = np.linalg.solve(ShS + lamda0 * np.eye(ShS.shape[0]), ShT).T - - return weights - - def apply_weights(self, kspace: np.ndarray, weights: Dict[int, np.ndarray]) -> np.ndarray: - """ - Applies the computed GRAPPA weights to the k-space data. - - Parameters: - ---------- - kspace : numpy.ndarray - The k-space data to apply the weights to. - - weights : dict - A dictionary containing the GRAPPA weights to apply. - - Returns: - ------- - numpy.ndarray: The reconstructed data after applying the weights. - """ - - # fin_shape = kspace.shape[:] - - # Put the coil dimension at the end - kspace = np.moveaxis(kspace, self.coil_axis, -1) - - # Get shape of kernel - kx, ky = self.kernel_size[:] - kx2, ky2 = int(kx / 2), int(ky / 2) - - # adjustment factor for odd kernel size - adjx = np.mod(kx, 2) - adjy = np.mod(ky, 2) - - # Pad kspace data - kspace = np.pad(kspace, ((kx2, kx2), (ky2, ky2), (0, 0)), mode="constant") - - with NTF() as frecon: - # Initialize recon array - recon = np.memmap(frecon, dtype=kspace.dtype, mode="w+", shape=kspace.shape) - - for ii in self.kernel_var_dict["patch_indices"]: - for xx, yy in zip( - self.kernel_var_dict["holes_x"][ii], self.kernel_var_dict["holes_y"][ii] - ): - # Collect sources for this hole and apply weights - S = kspace[xx - kx2 : xx + kx2 + adjx, yy - ky2 : yy + ky2 + adjy, :] - S = S[self.kernel_var_dict["patches"][ii, ...]] - recon[xx, yy, :] = (weights[ii] @ S[:, None]).squeeze() - - return np.moveaxis((recon[:] + kspace)[kx2:-kx2, ky2:-ky2, :], -1, self.coil_axis) - - -def et_query( - root: etree.Element, qlist: Sequence[str], namespace: str = "http://www.ismrm.org/ISMRMRD" -) -> str: - """ - ElementTree query function. - - This function queries an XML document using ElementTree. - - Parameters: - ----------- - root : Element - Root of the XML document to search through. - qlist : Sequence of str - A sequence of strings for nested searches, e.g., ["Encoding", "matrixSize"]. - namespace : str, optional - XML namespace to prepend query. - - Returns: - -------- - str - The retrieved data as a string. - """ - s = "." - prefix = "ismrmrd_namespace" - - ns = {prefix: namespace} - - for el in qlist: - s = s + f"//{prefix}:{el}" - - value = root.find(s, ns) - if value is None: - raise RuntimeError("Element not found") - - return str(value.text) - - -def get_padding(hdr: str) -> float: - """ - Extract the padding value from an XML header string. - - Parameters: - ----------- - hdr : str - The XML header string. - - Returns: - -------- - float - The padding value calculated as (x - max_enc)/2, where x is the readout dimension and - max_enc is the maximum phase-encoding dimension. - """ - et_root = etree.fromstring(hdr) - lims = ["encoding", "encodingLimits", "kspace_encoding_step_1"] - enc_limits_max = int(et_query(et_root, lims + ["maximum"])) + 1 - enc = ["encoding", "encodedSpace", "matrixSize"] - enc_x = int(et_query(et_root, enc + ["x"])) - padding = (enc_x - enc_limits_max) / 2 - - print(f"Padding from header: {padding}") - return padding - - -def get_padding_from_image_recon(k_space_shape: np.ndarray, image_recon_shape: np.ndarray) -> float: - """ - Calculate the padding value based on the shapes of the k-space data and the reconstructed image. - - Parameters: - ----------- - k_space_shape : np.ndarray - The shape of the k-space data, typically in the format (num_avg, num_slices, num_coils, num_ro, num_pe). - image_recon_shape : np.ndarray - The shape of the reconstructed image, typically in the format (num_avg, num_slices, num_coils, num_x, num_y). - - Returns: - -------- - float - The padding value calculated as (x - max_enc)/2, where x is the readout dimension from the k-space shape and - max_enc is the maximum phase-encoding dimension from the reconstructed image shape. - """ - enc_limits_max = image_recon_shape[ - -1 - ] # Assuming last dimension corresponds to phase-encoding direction - enc_x = k_space_shape[-2] # Assuming second to last dimension corresponds to readout direction - padding = (enc_x - enc_limits_max) / 2 - - print(f"Padding from image recon shapes: {padding}") - return padding - - -def zero_pad_kspace_slice_hdr( - hdr: str, unpadded_kspace: np.ndarray, image_recon_shape: np.ndarray -) -> np.ndarray: - """ - Perform zero-padding on k-space data to have the same number of - points in the x- and y-directions. - - Parameters - ---------- - hdr : str - The XML header string. - unpadded_kspace : array-like of shape (ro , coils, pe) - The k-space data to be padded. - - Returns - ------- - padded_kspace : ndarray of shape (ro_padded, coils, pe_padded) - The zero-padded k-space data, where ro_padded and pe_padded are - the dimensions of the readout and phase-encoding directions after - padding. - - Notes - ----- - The padding value is calculated using the `get_padding` function, which - extracts the padding value from the XML header string. If the difference - between the readout dimension and the maximum phase-encoding dimension - is not divisible by 2, the padding is applied asymmetrically, with one - side having an additional zero-padding. - - """ - padding = get_padding(hdr) - print(f"Calculated padding: {padding}") - padding2 = get_padding_from_image_recon(unpadded_kspace.shape, image_recon_shape) - print(f"Calculated padding from image recon shapes ({image_recon_shape}): {padding2}") - if padding % 2 != 0: - padding_left = int(np.floor(padding)) - padding_right = int(np.ceil(padding)) - else: - padding_left = int(padding) - padding_right = int(padding) - padded_kspace = np.pad(unpadded_kspace, ((0, 0), (0, 0), (padding_left, padding_right))) - - return padded_kspace - class FastMRIProstateDataset(torch.utils.data.Dataset): - def __init__(self, data_path: str, num_samples: None | int = None) -> None: + def __init__(self, data_path: str, num_samples: None | int = None, slice_index: str="middle") -> None: self.data_path = data_path self.num_samples = num_samples - self.kspace_data, self.image_data = self.pre_calc_kspace_with_grappa() + self.slice_index = slice_index + self.image_data = self.get_image_data() if num_samples is not None: - self.kspace_data = self.kspace_data[:num_samples] self.image_data = self.image_data[:num_samples] else: - self.num_samples = len(self.kspace_data) + self.num_samples = len(self.image_data) - def pre_calc_kspace_with_grappa(self) -> np.ndarray: - kspace_result_list = [] + def get_image_data(self) -> np.ndarray: image_result_list = [] for sample_idx, filename in enumerate(glob.glob(os.path.join(self.data_path, "*.h5"))): if (self.num_samples is not None) and (sample_idx >= self.num_samples): break try: with h5py.File(filename, "r") as hf: - kspace_data = hf["kspace"][:] - calibration_data = hf["calibration_data"][:] - hdr = hf["ismrmrd_header"][()] image_recon = hf["reconstruction_rss"][:] except Exception as e: print(f"Error processing file {filename}: {e}") continue - # kspace data: (num_avg, num_slices, num_coils, num_ro, num_pe) - # Calib_data: (num_slices, num_coils, num_pe_cal) - - # middle slice: - kspace_middle_slice = kspace_data[:, kspace_data.shape[1] // 2] - imwrite( - "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/kspace_middle_slice.tiff", - kspace_middle_slice, - ) - cal_middle_slice = calibration_data[calibration_data.shape[0] // 2] - print( - f"Processing file {filename} with kspace shape {kspace_middle_slice.shape} and calib shape {cal_middle_slice.shape}" - ) - imwrite( - "/home/melanie.dohmen/mri_recon/reports/test_prostate_T2_recon/calibration_middle_slice.tiff", - cal_middle_slice, - ) - - ####### - - kspace_slice_regridded = kspace_data[0, 0] - grappa_obj = Grappa( - np.transpose(kspace_slice_regridded, (2, 0, 1)), kernel_size=(5, 5), coil_axis=1 - ) - - kspace_slice_regridded_2 = kspace_data[1, 0] - grappa_obj_2 = Grappa( - np.transpose(kspace_slice_regridded_2, (2, 0, 1)), kernel_size=(5, 5), coil_axis=1 - ) - - # calculate GRAPPA weights for middle slice: - - calibration_regridded = calibration_data[calibration_data.shape[0] // 2, ...] - grappa_weight = grappa_obj.compute_weights( - np.transpose(calibration_regridded, (2, 0, 1)) - ) - grappa_weight2 = grappa_obj_2.compute_weights( - np.transpose(calibration_regridded, (2, 0, 1)) - ) - - #### - kspace_post_grappa_slice = np.zeros(shape=kspace_middle_slice.shape, dtype=complex) - kspace_post_grappa_slice_padded: Dict[int, np.ndarray] = {} - for average, grappa_obj, grappa_weight_dict in zip( - [0, 1, 2], - [grappa_obj, grappa_obj_2, grappa_obj], - [grappa_weight, grappa_weight2, grappa_weight], - ): - kspace_slice_regridded = kspace_middle_slice[average, ...] - kspace_post_grappa = grappa_obj.apply_weights( - np.transpose(kspace_slice_regridded, (2, 0, 1)), grappa_weight_dict - ) - kspace_post_grappa_slice[average] = np.moveaxis( - np.moveaxis(kspace_post_grappa, 0, 1), 1, 2 - ) - - # pad: - kspace_post_grappa_slice_padded[average] = zero_pad_kspace_slice_hdr( - hdr, kspace_post_grappa_slice[average], image_recon.shape - ) + if self.slice_index == "middle": + image_result_list.append(image_recon[image_recon.shape[0]//2]) + else: + image_result_list.extend([image_recon[i, :, :] for i in range(image_recon.shape[0])]) - # stack k-space data for all averages: - kspace_result_list.append( - np.stack(list(kspace_post_grappa_slice_padded.values()), axis=0) - ) - image_result_list.append(image_recon) - return kspace_result_list, image_result_list + return image_result_list def __len__(self) -> int: - return len(self.kspace_data) + return len(self.image_data) def __getitem__(self, idx: int) -> torch.Tensor: - kspace = self.kspace_data[idx] - return torch.from_numpy(kspace), torch.from_numpy(self.image_data[idx]) + # add batch dimension and convert to torch.Tensor + return torch.from_numpy(self.image_data[idx]).unsqueeze(0) From 58a2800a502dbb18234c2e9529e3937106aea0c0 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Mon, 1 Jun 2026 06:08:38 +0000 Subject: [PATCH 03/22] correct distortion names, remove debugging prints --- examples/config.yaml | 6 +++--- examples/run_all.py | 3 --- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/examples/config.yaml b/examples/config.yaml index eaa8040..673a0a9 100644 --- a/examples/config.yaml +++ b/examples/config.yaml @@ -13,19 +13,19 @@ distortions: - "CartesianUndersamplingEquispaced" - "CartesianUndersamplingEquispacedZeroACS" - "PartialFourier" - - "Phase-EncodeGhosting" + - "PhaseEncodeGhosting" - "SegmentedTranslationMotion" - "SegmentedRotationalMotion" - "TranslationMotion" - "RotationalMotion" - - "Off-centerAnisotropicGaussianBiasField" + - "OffCenterAnisotropicGaussianBiasField" - "GaussianBiasField" - "AnisotropicLP" - "HannTaperLP" - "KaiserTaperLP" - "GaussianNoise" - "IsotropicLP" - - "RadialHigh-passEmphasis" + - "RadialHighPassEmphasis" reconstruction_algorithms: - "zero-filled" - "conjugate-gradient" diff --git a/examples/run_all.py b/examples/run_all.py index dbb4f3a..1473fde 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -22,7 +22,6 @@ choose_distortion, ) from mri_recon.reconstruction import ( - ConjugateGradientReconstructor, choose_reconstructor, uses_oasis_centered_path, compatible_dataset_with_reconstructor, @@ -31,7 +30,6 @@ OasisCenteredFFTPhysics, OasisCenterSliceFolderDataset, FastMRIProstateDataset, - fastmri_measurement_to_image, fastmri_measurement_to_oasis_kspace, oasis_kspace_to_fastmri_measurement, image_to_kspace, @@ -69,7 +67,6 @@ def get_measurement_sample( y_centered = fastmri_measurement_to_oasis_kspace(y, device=run_device) # reconstructed reference image: # shape: (B, 1, H, W) dtype: float32 - print("stop for testing") elif dataset_name == "fastmri_brain": # reference image, shape: (B, 1, H/2, H/2) dtype: float32 x = sample_batch[0].to(run_device) From 2969f9d0d6203381b9704c717a9577420a58f678 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Mon, 1 Jun 2026 10:10:57 +0000 Subject: [PATCH 04/22] correct coil dimensions and improve run_all script --- examples/config_md.yaml | 63 ++++++++++++++++---------------- examples/run_all.py | 41 +++++++++++---------- mri_recon/distortions/base.py | 28 +++----------- mri_recon/utils/__init__.py | 2 + mri_recon/utils/oasis_adapter.py | 35 +++++++----------- mri_recon/utils/plot.py | 2 +- 6 files changed, 75 insertions(+), 96 deletions(-) diff --git a/examples/config_md.yaml b/examples/config_md.yaml index ec94872..a85bc17 100644 --- a/examples/config_md.yaml +++ b/examples/config_md.yaml @@ -1,44 +1,43 @@ data: - "fastmri_knee": "/home/melanie.dohmen/ArtifactLab/data/singlecoil_val" - "oasis": "/home/melanie.dohmen/ArtifactLab/data/oasis" - "fastmri_brain": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_multicoil_brain_test" - "cmrxrecon": "/home/melanie.dohmen/ArtifactLab/data/CMRxRecon" - "fastmri_prostate": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_prostate_T2_IDS_001_020" - + #"fastmri_knee": "/home/melanie.dohmen/ArtifactLab/data/singlecoil_val" + "oasis": "/home/melanie.dohmen/ArtifactLab/data/oasis" + "fastmri_brain": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_multicoil_brain_test" + "cmrxrecon": "/home/melanie.dohmen/ArtifactLab/data/CMRxRecon" + "fastmri_prostate": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_prostate_T2_IDS_001_020" distortions: - #- "BaseDistortion" + - "BaseDistortion" #- "CartesianUndersamplingVariableDensity" - #- "Cartesian undersampling (uniform random)", - #- "Cartesian undersampling (uniform random, zero ACS)", - #- "Cartesian undersampling (equispaced)", + #- "CartesianUndersamplingUniformRandom" + #- "CartesianUndersamplingUniformRandomZeroACS" + #- "CartesianUndersamplingEquispaced" #- "CartesianUndersamplingEquispacedZeroACS" - #- "Partial Fourier", - #- "Phase-encode ghosting", - #- "Segmented translation motion", - # "Segmented rotational motion", - # "Translation motion", - # "Rotational motion", - # "Off-center anisotropic Gaussian bias field", - # "Gaussian bias field", - # "Anisotropic LP", - # "Hann taper LP", - # "Kaiser taper LP", - - "GaussianNoise" - # "Isotropic LP", - # "Radial high-pass emphasis", + #- "PartialFourier" + #- "PhaseEncodeGhosting" + #- "SegmentedTranslationMotion" + #- "SegmentedRotationalMotion" + #- "TranslationMotion" + #- "RotationalMotion" + #- "OffCenterAnisotropicGaussianBiasField" + #- "GaussianBiasField" + #- "AnisotropicLP" + #- "HannTaperLP" + #- "KaiserTaperLP" + #- "GaussianNoise" + #- "IsotropicLP" + #- "RadialHighPassEmphasis" reconstruction_algorithms: - "zero-filled" - "conjugate-gradient" - #- "ram", - #- "dip", - #- "tv-pgd", - #- "wavelet-fista", - #- "tv-fista", - #- "tv-pdhg", + - "ram" + - "dip" + - "tv-pgd" + - "wavelet-fista" + - "tv-fista" + - "tv-pdhg" - "unet-fastmri" - "unet-oasis-acceleration4" - #- "unet-oasis-acceleration8" - #- "unet-oasis-acceleration10" + - "unet-oasis-acceleration8" + - "unet-oasis-acceleration10" num_samples: 1 keep_fraction: 0.25 center_fraction: 0.125 diff --git a/examples/run_all.py b/examples/run_all.py index 1473fde..f0d61a0 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -1,7 +1,7 @@ """Inference various reconstructors for various distortion operators. Usage: - python examples/fastmri_inference_plot.py --source ../ram-experiments/data/fastmri/knee/singlecoil_val + python examples/run_all.py config.yaml """ import os @@ -10,7 +10,7 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - +from datetime import datetime import deepinv as dinv import torch import yaml @@ -57,7 +57,7 @@ def get_measurement_sample( # centered k-space data, shape: (B, 2, H, W) dtype: float32 y_centered = image_to_kspace(x) # k-space data, shape: (B, 2, H, W) dtype: float32 - y = oasis_kspace_to_fastmri_measurement(y_centered, device=run_device) + y = oasis_kspace_to_fastmri_measurement(y_centered) elif dataset_name == "fastmri_knee": # reference image, shape: (B, 1, H/2, H/2) dtype: float32 x = sample_batch[0].to(run_device) @@ -124,7 +124,7 @@ def get_measurement_sample( # create oasis-like k-space data from image: y_centered = image_to_kspace(x) print(f"\t[Debug] Centered k-space shape: {y_centered.shape}, type: {y_centered.dtype}") - y = oasis_kspace_to_fastmri_measurement(y_centered, device=run_device) + y = oasis_kspace_to_fastmri_measurement(y_centered) print(f"\tk-space shape {y.shape}[{y.dtype}] and reference image shape: {x.shape}[{x.dtype}]") if coil_maps is not None: @@ -225,11 +225,11 @@ def get_measurement_sample( ) for reconstructor_name in config["reconstruction_algorithms"]: - print(f"\t\t{reconstructor_name} ...") - if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): - - # only run on reconstructors, that use the fastmri-like k-space - if not uses_oasis_centered_path(dataset_name, reconstructor_name): + # only run on reconstructors, that use the fastmri-like k-space + if not uses_oasis_centered_path(dataset_name, reconstructor_name): + print(f"\t\t{reconstructor_name} ...") + start = datetime.now() + if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): reconstructor = choose_reconstructor( reconstructor_name, @@ -265,7 +265,7 @@ def get_measurement_sample( x_uncorrected = physics_clean.crop(x_uncorrected, shape=x_reference.shape[-2:]) if x_corrected.shape[-2:] != x_reference.shape[-2:]: - x_corrected_clean = physics_distorted.crop(x_corrected, shape=x_reference.shape[-2:]) + x_corrected = physics_distorted.crop(x_corrected, shape=x_reference.shape[-2:]) # save reconstructed images imwrite( @@ -282,6 +282,7 @@ def get_measurement_sample( ), convert_image_for_save(x_corrected), ) + print(f"\t\t... done in {start- datetime.now()}") except Exception as e: print( @@ -289,8 +290,8 @@ def get_measurement_sample( ) - else: - print(f"\t\t ... not compatible with {dataset_name}") + else: + print(f"\t\t ... not compatible with {dataset_name}") # now proceed with oasis-centered fft path @@ -313,11 +314,11 @@ def get_measurement_sample( ) for reconstructor_name in config["reconstruction_algorithms"]: - print(f"\t\t{reconstructor_name} ...") - if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): - - # skip all reconstructors, that don't use the oasis-centered path - if uses_oasis_centered_path(dataset_name, reconstructor_name): + + # skip all reconstructors, that don't use the oasis-centered path + if uses_oasis_centered_path(dataset_name, reconstructor_name): + print(f"\t\t{reconstructor_name} ...") + if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): reconstructor = choose_reconstructor( reconstructor_name, @@ -351,7 +352,7 @@ def get_measurement_sample( x_uncorrected = physics_clean.crop(x_uncorrected, shape=x_reference.shape[-2:]) if x_corrected.shape[-2:] != x_reference.shape[-2:]: - x_corrected_clean = physics_distorted.crop(x_corrected, shape=x_reference.shape[-2:]) + x_corrected = physics_distorted.crop(x_corrected, shape=x_reference.shape[-2:]) # save reconstructed images @@ -375,6 +376,6 @@ def get_measurement_sample( f"\t\tError using {reconstructor_name} with distortion {distortion_name} on sample {i}: {e}" ) - else: - print(f"\t\t ... not compatible with {dataset_name}") + else: + print(f"\t\t ... not compatible with {dataset_name}") diff --git a/mri_recon/distortions/base.py b/mri_recon/distortions/base.py index 51d9e0a..4178a17 100644 --- a/mri_recon/distortions/base.py +++ b/mri_recon/distortions/base.py @@ -176,35 +176,19 @@ def __init__(self, distortion: BaseDistortion = None, *args, **kwargs): def A(self, x: torch.Tensor) -> torch.Tensor: y = super().A(x) - y = y.squeeze(2) # remove coil dim if singlecoil + + y = y.squeeze(2) # remove coil dimension if single coil + return self.distortion(y) def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: - + if len(y.shape) == (5 if self.three_d else 4): - y = y.unsqueeze(2) # add coil dim if singlecoil + y = y.unsqueeze(2) # add coil dimension for single coil y = self.distortion.A_adjoint(y) - # # in order to match the reference image shape - # # a crop must be performed here on k-space data: - # if y.shape[-2:] != self.img_size[-2:]: - # y =self.crop(y, crop=True) - - # # and on coil maps for multi-coil data: - # if self.coil_maps is not None and self.coil_maps.shape[2:] != self.img_size[1:]: - # self.coil_maps = self.crop(self.coil_maps, crop=True) return super().A_adjoint(y) - # def A_dagger(self, y: torch.Tensor, **kwargs) -> torch.Tensor: - - # # in order to match the reference image shape - # # a crop must be performed here on k-space data: - # if y.shape[-2:] != self.img_size[-2:]: - # y =self.crop(y, crop=True) - - # # and on coil maps for multi-coil data: - # if self.coil_maps is not None and self.coil_maps.shape[2:] != self.img_size[1:]: - # self.coil_maps = self.crop(self.coil_maps, crop=True) - + # def A_dagger(self, y: torch.Tensor, **kwargs) -> torch.Tensor: # return super().A_dagger(y, coil_maps=self.coil_maps, **kwargs) diff --git a/mri_recon/utils/__init__.py b/mri_recon/utils/__init__.py index ed72870..0436aa5 100644 --- a/mri_recon/utils/__init__.py +++ b/mri_recon/utils/__init__.py @@ -14,6 +14,8 @@ ) from .oasis_adapter import image_to_kspace as image_to_kspace from .oasis_adapter import kspace_to_image as kspace_to_image +from .oasis_adapter import image_to_fastmri_measurement as image_to_fastmri_measurement +from .oasis_adapter import oasis_kspace_to_fastmri_measurement as oasis_kspace_to_fastmri_measurement from .prostate_adaptor import FastMRIProstateDataset as FastMRIProstateDataset from .plot import save_kspace_plot as save_kspace_plot from .plot import _kspace_to_log_magnitude as _kspace_to_log_magnitude diff --git a/mri_recon/utils/oasis_adapter.py b/mri_recon/utils/oasis_adapter.py index 89f3ce6..f9542b8 100644 --- a/mri_recon/utils/oasis_adapter.py +++ b/mri_recon/utils/oasis_adapter.py @@ -296,15 +296,11 @@ def fastmri_measurement_to_oasis_kspace( Centered OASIS-convention k-space tensor with shape ``(B, 2, H, W)``. """ - result = image_to_kspace(fastmri_measurement_to_image(y, coil_maps=coil_maps, device=device)) + return image_to_kspace(fastmri_measurement_to_image(y, coil_maps=coil_maps, device=device)) - - return result - -def image_to_fast_mri_measurement( +def image_to_fastmri_measurement( x: torch.Tensor, - coil_maps: torch.Tensor | None = None, device: torch.device | str | None = None, ) -> torch.Tensor: """Perform FFT from image space to k-space (fast-MRI convention). @@ -313,8 +309,6 @@ def image_to_fast_mri_measurement( ---------- x : torch.Tensor image tensor with shape ``(B, 2, H, W)``. - coil_maps : torch.Tensor | None, optional - Coil sensitivity maps with shape ``(B, C, H, W)``, where ``C`` is the number of coils. device : torch.device | str, optional Device on which to instantiate the temporary native physics operator. @@ -326,10 +320,9 @@ def image_to_fast_mri_measurement( if device is None: device = x.device - physics = DistortedKspaceMultiCoilMRI( - distortion=BaseDistortion(), + physics = dinv.physics.MultiCoilMRI( img_size=(1, 2, *x.shape[-2:]), - coil_maps=coil_maps, + coil_maps=None, device=device, ) return physics.A(x) @@ -352,10 +345,10 @@ def oasis_kspace_to_fastmri_measurement( FastMRI-convention k-space tensor with shape ``(B, 2, H, W)``. """ - return image_to_fast_mri_measurement(kspace_to_image(y)) + return image_to_fastmri_measurement(kspace_to_image(y)) -class OasisCenteredFFTPhysics(dinv.physics.LinearPhysics): +class OasisCenteredFFTPhysics(dinv.utils.mixins.MRIMixin, dinv.physics.LinearPhysics): """Physics adapter matching the OASIS U-Net FFT convention. Parameters @@ -400,13 +393,13 @@ def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: return kspace_to_image(self.distortion.A_adjoint(y)) - def A_dagger(self, y: torch.Tensor, **kwargs) -> torch.Tensor: - r""" - Computes least squares solution to the MRI inverse problem, as proposed in `SENSE: Sensitivity encoding for fast MRI `_. + # def A_dagger(self, y: torch.Tensor, **kwargs) -> torch.Tensor: + # r""" + # Computes least squares solution to the MRI inverse problem, as proposed in `SENSE: Sensitivity encoding for fast MRI `_. - By default uses conjugate gradient solver. Overwrite default solver arguments by passing `kwargs`. See :func:`deepinv.optim.linear.least_squares` for details. + # By default uses conjugate gradient solver. Overwrite default solver arguments by passing `kwargs`. See :func:`deepinv.optim.linear.least_squares` for details. - :param dict kwargs: kwargs to pass to base :meth:`deepinv.physics.LinearPhysics.A_dagger`. - :returns: (:class:`torch.Tensor`) image with shape `(B,2,...,H,W)` - """ - return super().A_dagger(y, **kwargs) + # :param dict kwargs: kwargs to pass to base :meth:`deepinv.physics.LinearPhysics.A_dagger`. + # :returns: (:class:`torch.Tensor`) image with shape `(B,2,...,H,W)` + # """ + # return super().A_dagger(y, **kwargs) diff --git a/mri_recon/utils/plot.py b/mri_recon/utils/plot.py index 5379519..4f18b1e 100644 --- a/mri_recon/utils/plot.py +++ b/mri_recon/utils/plot.py @@ -89,4 +89,4 @@ def convert_image_for_save(im: torch.Tensor) -> np.ndarray: """ if torch.is_complex(im) or im.shape[1] == 2: im = dinv.utils.signals.complex_abs(im, dim=1, keepdim=False) - return im.numpy() + return im.detach().cpu().numpy() From 26f720127351dbf0ee55147e81b02b7392fc5f2b Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Mon, 1 Jun 2026 10:16:43 +0000 Subject: [PATCH 05/22] ignore downloaded checkpoints and additional personalized config files --- .gitignore | 6 + examples/config.yaml | 2 +- examples/run_all.py | 81 +++-- mri_recon/distortions/__init__.py | 2 +- mri_recon/distortions/base.py | 7 +- .../examples/run_experiments_fastmri_brain.py | 313 ------------------ .../run_experiments_fastmri_prostateT2.py | 289 ---------------- mri_recon/reconstruction/inference.py | 8 +- mri_recon/utils/__init__.py | 1 - mri_recon/utils/oasis_adapter.py | 2 +- mri_recon/utils/plot.py | 2 +- mri_recon/utils/prostate_adaptor.py | 11 +- 12 files changed, 65 insertions(+), 659 deletions(-) delete mode 100644 mri_recon/examples/run_experiments_fastmri_brain.py delete mode 100644 mri_recon/examples/run_experiments_fastmri_prostateT2.py diff --git a/.gitignore b/.gitignore index 15c7e8d..80da5e0 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,9 @@ data/ *_cache/ reports/ .python-version +*.ckpt +*.pt +manifest.json +# Ignore all yaml files in the examples folder except config.yaml as template: +examples/*.yaml +!examples/config.yaml diff --git a/examples/config.yaml b/examples/config.yaml index 673a0a9..7fbed4e 100644 --- a/examples/config.yaml +++ b/examples/config.yaml @@ -43,4 +43,4 @@ num_samples: 1 keep_fraction: 0.25 center_fraction: 0.125 verbose: true -results_dir: "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1" \ No newline at end of file +results_dir: "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1" diff --git a/examples/run_all.py b/examples/run_all.py index f0d61a0..0141a68 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -38,7 +38,6 @@ ) - def get_measurement_sample( sample_batch: object, dataset_name: str, @@ -48,7 +47,7 @@ def get_measurement_sample( Always prepare a (fast-MRI-like) non-centered k-space measurement as well as a (oasis-like) centered k-space version of the measurement - and a reference reconstruction in the image domain. + and a reference reconstruction in the image domain. """ coil_maps = None if dataset_name == "oasis": @@ -82,7 +81,7 @@ def get_measurement_sample( ) # centered k-space data, shape: (B, 2, H, W) dtype: float32 y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) - + elif dataset_name == "cmrxrecon": # reference image, shape: (B, 2, n_timepoints, (n_coils), H, W) x = sample_batch[0].to(run_device) @@ -91,7 +90,7 @@ def get_measurement_sample( y = sample_batch[1].to(run_device) print(f"\t[Debug] k-space shape: {y.shape}, dtype: {y.dtype}") - # not available for all samples, either None or + # not available for all samples, either None or # shape (1, num_coils, H, W) coil_maps = ( sample_batch[2]["coil_maps"].to(run_device) @@ -100,15 +99,17 @@ def get_measurement_sample( and "coil_maps" in sample_batch[2] else None ) - print(f"\t[Debug] Coil maps shape: {coil_maps.shape if coil_maps is not None else None}, dtype: {coil_maps.dtype if coil_maps is not None else None}") + print( + f"\t[Debug] Coil maps shape: {coil_maps.shape if coil_maps is not None else None}, dtype: {coil_maps.dtype if coil_maps is not None else None}" + ) # centered k-space data, shape: (B, 2, H, W) dtype: float32 y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) print(f"\t[Debug] Centered k-space shape: {y_centered.shape}, dtype: {y_centered.dtype}") # reconstruct coil-combined image reference from multi-coil k-space data using # integrated espirit sensitivity map estimation, RSS coil combination - - #x = fastmri_measurement_to_image(y, coil_maps=coil_maps, rss=True) - #print(f"\t[Debug] Reference image shape: {x.shape}, dtype: {x.dtype}") + + # x = fastmri_measurement_to_image(y, coil_maps=coil_maps, rss=True) + # print(f"\t[Debug] Reference image shape: {x.shape}, dtype: {x.dtype}") elif dataset_name == "fastmri_prostate": # reference image, shape: (slices, W, H): dtype float32 x = sample_batch[0].to(run_device) @@ -117,9 +118,7 @@ def get_measurement_sample( # add zero imaginary channel: # (B, slices, H, W) -> (B, 2, slices, H, W) x = torch.stack([x, torch.zeros_like(x)], dim=1) - - # y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) # create oasis-like k-space data from image: y_centered = image_to_kspace(x) @@ -134,7 +133,6 @@ def get_measurement_sample( if __name__ == "__main__": - # read config file in yaml format as first argument from commmand line if len(sys.argv) < 2: print("Usage: python examples/run_all.py ") @@ -149,7 +147,6 @@ def get_measurement_sample( device = dinv.utils.get_device() for dataset_name, dataset_rootdir in config["data"].items(): - print(f"=== {dataset_name} ===") # initialize dataset @@ -175,13 +172,14 @@ def get_measurement_sample( apply_mask=False, ) elif dataset_name == "fastmri_prostate": - dataset = FastMRIProstateDataset(data_path=dataset_rootdir, num_samples=config["num_samples"], slice_index="middle") + dataset = FastMRIProstateDataset( + data_path=dataset_rootdir, num_samples=config["num_samples"], slice_index="middle" + ) else: raise NotImplementedError(f"Invalid dataset: {dataset_name}") # loop through samples of dataset for i, batch in enumerate(iter(torch.utils.data.DataLoader(dataset))): - # exit loop if we have processed the specified number of samples if i >= config["num_samples"]: break @@ -195,7 +193,9 @@ def get_measurement_sample( # save reference image imwrite( - os.path.join(config["results_dir"], f"image_{dataset_name}_sample_{i}_reference.tiff"), + os.path.join( + config["results_dir"], f"image_{dataset_name}_sample_{i}_reference.tiff" + ), convert_image_for_save(x_reference), ) @@ -203,7 +203,7 @@ def get_measurement_sample( physics_clean = DistortedKspaceMultiCoilMRI( BaseDistortion(), img_size=y.shape[-2:], coil_maps=coil_maps, device=device ) - + # reference from dataset: for distortion_name in config["distortions"]: print(f"\t{distortion_name} ...") @@ -229,8 +229,7 @@ def get_measurement_sample( if not uses_oasis_centered_path(dataset_name, reconstructor_name): print(f"\t\t{reconstructor_name} ...") start = datetime.now() - if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): - + if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): reconstructor = choose_reconstructor( reconstructor_name, img_size=y_distorted.shape[-2:], @@ -241,7 +240,8 @@ def get_measurement_sample( # save reference and distorted k-space for debugging purposes imwrite( os.path.join( - config["results_dir"], f"kspace_{dataset_name}_sample_{i}_reference.tiff" + config["results_dir"], + f"kspace_{dataset_name}_sample_{i}_reference.tiff", ), _kspace_to_log_magnitude(y).numpy(), ) @@ -255,17 +255,19 @@ def get_measurement_sample( # actual reconstruction with the selected reconstructor try: - x_uncorrected = reconstructor(y_distorted, physics_clean) x_corrected = reconstructor(y_distorted, physics_distorted) - # crop recostructed image to reference image size: if x_uncorrected.shape[-2:] != x_reference.shape[-2:]: - x_uncorrected = physics_clean.crop(x_uncorrected, shape=x_reference.shape[-2:]) + x_uncorrected = physics_clean.crop( + x_uncorrected, shape=x_reference.shape[-2:] + ) if x_corrected.shape[-2:] != x_reference.shape[-2:]: - x_corrected = physics_distorted.crop(x_corrected, shape=x_reference.shape[-2:]) + x_corrected = physics_distorted.crop( + x_corrected, shape=x_reference.shape[-2:] + ) # save reconstructed images imwrite( @@ -282,18 +284,14 @@ def get_measurement_sample( ), convert_image_for_save(x_corrected), ) - print(f"\t\t... done in {start- datetime.now()}") + print(f"\t\t... done in {start - datetime.now()}") except Exception as e: - print( - f"Error using {reconstructor_name}: {e}" - ) + print(f"Error using {reconstructor_name}: {e}") - else: print(f"\t\t ... not compatible with {dataset_name}") - # now proceed with oasis-centered fft path physics_clean = OasisCenteredFFTPhysics(BaseDistortion()) @@ -306,20 +304,17 @@ def get_measurement_sample( cartesian_axis=-1, ) - - y_distorted = torch.fft.fftshift(distortion.A(torch.fft.fftshift(y_centered, dim=(-1, -2))), dim=(-2, -1)) - - physics_distorted = OasisCenteredFFTPhysics( - distortion + y_distorted = torch.fft.fftshift( + distortion.A(torch.fft.fftshift(y_centered, dim=(-1, -2))), dim=(-2, -1) ) + physics_distorted = OasisCenteredFFTPhysics(distortion) + for reconstructor_name in config["reconstruction_algorithms"]: - # skip all reconstructors, that don't use the oasis-centered path if uses_oasis_centered_path(dataset_name, reconstructor_name): print(f"\t\t{reconstructor_name} ...") if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): - reconstructor = choose_reconstructor( reconstructor_name, img_size=y_distorted.shape[-2:], @@ -330,7 +325,8 @@ def get_measurement_sample( # save reference and distorted k-space for debugging purposes imwrite( os.path.join( - config["results_dir"], f"kspace_centered_{dataset_name}_sample_{i}_reference.tiff" + config["results_dir"], + f"kspace_centered_{dataset_name}_sample_{i}_reference.tiff", ), _kspace_to_log_magnitude(y_centered).numpy(), ) @@ -344,16 +340,18 @@ def get_measurement_sample( # actual reconstruction with the algo being evaluated try: - x_uncorrected = reconstructor(y_distorted, physics_clean) x_corrected = reconstructor(y_distorted, physics_distorted) if x_uncorrected.shape[-2:] != x_reference.shape[-2:]: - x_uncorrected = physics_clean.crop(x_uncorrected, shape=x_reference.shape[-2:]) + x_uncorrected = physics_clean.crop( + x_uncorrected, shape=x_reference.shape[-2:] + ) if x_corrected.shape[-2:] != x_reference.shape[-2:]: - x_corrected = physics_distorted.crop(x_corrected, shape=x_reference.shape[-2:]) - + x_corrected = physics_distorted.crop( + x_corrected, shape=x_reference.shape[-2:] + ) # save reconstructed images imwrite( @@ -378,4 +376,3 @@ def get_measurement_sample( else: print(f"\t\t ... not compatible with {dataset_name}") - diff --git a/mri_recon/distortions/__init__.py b/mri_recon/distortions/__init__.py index 1b99916..ea85943 100644 --- a/mri_recon/distortions/__init__.py +++ b/mri_recon/distortions/__init__.py @@ -20,4 +20,4 @@ RadialHighPassEmphasisDistortion, ) from .undersampling import CartesianUndersampling, PartialFourierDistortion -from .utils import choose_distortion \ No newline at end of file +from .utils import choose_distortion diff --git a/mri_recon/distortions/base.py b/mri_recon/distortions/base.py index 4178a17..d89efd0 100644 --- a/mri_recon/distortions/base.py +++ b/mri_recon/distortions/base.py @@ -178,17 +178,16 @@ def A(self, x: torch.Tensor) -> torch.Tensor: y = super().A(x) y = y.squeeze(2) # remove coil dimension if single coil - + return self.distortion(y) def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: - if len(y.shape) == (5 if self.three_d else 4): y = y.unsqueeze(2) # add coil dimension for single coil y = self.distortion.A_adjoint(y) - + return super().A_adjoint(y) - # def A_dagger(self, y: torch.Tensor, **kwargs) -> torch.Tensor: + # def A_dagger(self, y: torch.Tensor, **kwargs) -> torch.Tensor: # return super().A_dagger(y, coil_maps=self.coil_maps, **kwargs) diff --git a/mri_recon/examples/run_experiments_fastmri_brain.py b/mri_recon/examples/run_experiments_fastmri_brain.py deleted file mode 100644 index e482e79..0000000 --- a/mri_recon/examples/run_experiments_fastmri_brain.py +++ /dev/null @@ -1,313 +0,0 @@ -from datetime import datetime -import os -import sys -import h5py -from tifffile import imwrite -import torch -import deepinv as dinv -import pandas as pd - - -sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) - -from mri_recon.distortions import DistortedKspaceMultiCoilMRI, BaseDistortion, choose_distortion -from mri_recon.reconstruction import choose_reconstructor -from mri_recon.utils.oasis_adapter import ( - DistortedOasisMeasurement, - fastmri_measurement_to_oasis_kspace, - kspace_to_image, -) - - -sensitivity_map_estimation_algorithm = [ - "espirit", - # "unity", - # "birdcage", -] - - -DISTORTIONS = [ - "BaseDistortion", - # "PhaseEncodeGhosting", - # "CartesianUndersamplingVariableDensity", - # "CartesianUndersamplingUniformRandom", - # "HannTaperLP", - # "KaiserTaperLP", - # "RadialHighPassEmphasis", - # "IsotropicLP", - # "OffCenterAnisotropicGaussianKspaceBiasField", - # "TranslationMotion", - # "RotationalMotion", - # "SegmentedRotationalMotion", - # "SegmentedTranslationMotion", - # "GaussianKspaceBiasField", - # "GaussianNoise", -] - -RECONSTRUCTORS = [ - # "zero-filled", - # "conjugate-gradient", - # "ram", - # "dip", - # "tv-pgd", - # "wavelet-fista", - # "tv-fista", - # "tv-pdhg", - # "unet", # will trigger download of pretrained weights if not already present - # *list(EXPLICIT_UNET_ALGORITHMS) - "unet-fastmri", - "unet-oasis-acceleration4", - # "unet-oasis-acceleration8", - # "unet-oasis-acceleration10", -] - -reference_reconstructor = "conjugate-gradient" -reference_map_estimation = "espirit" - -filenames = [ - "/home/melanie.dohmen/mri_recon/data/fastmri/multicoil_brain_test/file_brain_AXFLAIR_200_6002452.h5" -] - -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -result_path = "/home/melanie.dohmen/mri_recon/reports/experiments_fastmri_brain/" - -os.makedirs(result_path, exist_ok=True) - -distortion_times = {} -reconstruction_times = {} - -for f_idx, filename in enumerate(filenames): - with h5py.File(filename, "r") as hf: - kspace_data = hf["kspace"][:] - reconstruction_rss = hf["reconstruction_rss"][:] - # hdr = hf["ismrmrd_header"][()] - # print(hf.keys()) - - x = torch.from_numpy(reconstruction_rss).unsqueeze(0).unsqueeze(0) - y = torch.view_as_real(torch.from_numpy(kspace_data)).unsqueeze(0).moveaxis(-1, 1) - # image shape: (1, channels, slices, H, W) - print("image x.shape:", x.shape) - # k-space shape: (1, channels, slices, coils, H, W) - print("k-space y.shape:", y.shape) - x = x.to(device) - y = y.to(device) - - # select middle slice: - x = x[:, :, x.shape[2] // 2, ...] - y = y[:, :, y.shape[2] // 2, ...] - - # read size of dimensions: - batch_size, channels, n_coils, ksp_w, ksp_h = y.shape - - for map_estimation in sensitivity_map_estimation_algorithm: - print(f"Estimating coil maps with {map_estimation}...") - if map_estimation == "espirit": - estimate_coil_maps = dinv.datasets.MRISliceTransform( - estimate_coil_maps=True, - acs=15, # Num. low frequency, fix to 15 - ) - _, _, params = estimate_coil_maps(target=x[0], kspace=y[0]) - - coil_maps = params["coil_maps"] - print("estimated coil maps shape: ", coil_maps.shape) - - coil_maps_result_path = os.path.join( - result_path, f"brain_sample_{f_idx}_coil_maps_{map_estimation}" - ) - os.makedirs(coil_maps_result_path, exist_ok=True) - for c_idx in range(coil_maps.shape[0]): - imwrite( - os.path.join( - coil_maps_result_path, - f"brain_sample_{f_idx}_{map_estimation}_map_{c_idx}.tiff", - ), - coil_maps[c_idx].abs().numpy(), - ) - - elif map_estimation == "unity": - coil_maps = torch.ones((n_coils, ksp_w, ksp_h), dtype=torch.complex64, device=device) - print("estimated coil maps shape: ", coil_maps.shape) - coil_maps_result_path = os.path.join( - result_path, f"brain_sample_{f_idx}_coil_maps_{map_estimation}" - ) - os.makedirs(coil_maps_result_path, exist_ok=True) - for c_idx in range(coil_maps.shape[0]): - imwrite( - os.path.join( - coil_maps_result_path, - f"brain_sample_{f_idx}_{map_estimation}_map_{c_idx}.tiff", - ), - coil_maps[c_idx].abs().numpy(), - ) - - elif map_estimation == "birdcage": - coil_maps = n_coils - - if map_estimation == reference_map_estimation: - physics_clean = dinv.physics.MultiCoilMRI( - img_size=(ksp_w, ksp_h), - mask=None, - coil_maps=coil_maps, - device=device, - ) - - y_distorted = BaseDistortion()(y) - print("base distortion does not change y:", torch.all(y_distorted == y)) - print("kspace distorted.shape: ", y_distorted.shape) - - x_recon_reference = choose_reconstructor(reference_reconstructor)( - y_distorted, physics_clean - ) - print("reconstructed image shape: ", x_recon_reference.shape) - - x_recon_reference_cropped = physics_clean.crop(x_recon_reference, shape=x.shape) - print("cropped reconstructed image shape: ", x_recon_reference_cropped.shape) - - x_as_inversed_y = physics_clean.A_adjoint(y_distorted, rss=True) - print("inversed y shape: ", x_as_inversed_y.shape) - x_as_inversed_y_cropped = physics_clean.crop(x_as_inversed_y, shape=x.shape) - print("cropped inversed y shape: ", x_as_inversed_y_cropped.shape) - - imwrite( - os.path.join( - result_path, - f"brain_sample_{f_idx}_{map_estimation}_reconstructed_reference.tiff", - ), - x_recon_reference_cropped[0, 0].abs().numpy(), - ) - imwrite( - os.path.join(result_path, f"brain_sample_{f_idx}_{map_estimation}_inversed_y.tiff"), - x_as_inversed_y_cropped[0, 0].abs().numpy(), - ) - - if map_estimation == "birdcage": - coil_maps = physics_clean.coil_maps - print("estimated coil maps shape: ", coil_maps.shape) - coil_maps_result_path = os.path.join( - result_path, f"brain_sample_{f_idx}_coil_maps_{map_estimation}" - ) - os.makedirs(coil_maps_result_path, exist_ok=True) - for c_idx in range(coil_maps.shape[0]): - imwrite( - os.path.join( - coil_maps_result_path, - f"brain_sample_{f_idx}_{map_estimation}_map_{c_idx}.tiff", - ), - coil_maps[0, c_idx].abs().numpy(), - ) - - for distortion_name in DISTORTIONS: - print("Distortion: ", distortion_name) - - start = datetime.now() - - distortion = choose_distortion(distortion_name) - physics_distorted = DistortedKspaceMultiCoilMRI( - distortion=distortion, - img_size=(ksp_w, ksp_h), - mask=None, - coil_maps=coil_maps, - device=device, - ) - - print("inserting k-space into distortion with shape: ", y.shape) - y_distorted = distortion(y) - - x_distorted_as_inversed_y = physics_distorted.A_adjoint(y_distorted, rss=True) - x_distorted_as_inversed_y_cropped = physics_distorted.crop( - x_distorted_as_inversed_y, shape=x.shape - ) - imwrite( - os.path.join( - result_path, - f"brain_sample_{f_idx}_{map_estimation}_{distortion_name}_inversed_y.tiff", - ), - x_distorted_as_inversed_y_cropped[0, 0].abs().numpy(), - ) - - for reconstructor_name in RECONSTRUCTORS: - print("Reconstructor: ", reconstructor_name) - - start_recon = datetime.now() - - if reconstructor_name in [ - "unet-oasis-acceleration4", - "unet-oasis-acceleration8", - "unet-oasis-acceleration10", - ]: - physics_distorted = DistortedOasisMeasurement( - distortion=distortion, - img_size=(ksp_w, ksp_h), - mask=None, - coil_maps=coil_maps, - device=device, - ) - y_distorted_for_recon = kspace_to_image( - fastmri_measurement_to_oasis_kspace(y_distorted) - ) - else: - y_distorted_for_recon = y_distorted - - reconstructor = choose_reconstructor( - reconstructor_name, - img_size=y.shape[-2:], - device=device, - verbose=True, - ).to(device) - - try: - x_distorted = reconstructor(y_distorted_for_recon, physics_distorted) - - x_distorted_cropped = ( - physics_distorted.crop(x_distorted, shape=x.shape).detach().cpu() - ) - - imwrite( - os.path.join( - result_path, - f"brain_sample_{f_idx}_{map_estimation}_{distortion_name}_{reconstructor_name}.tiff", - ), - physics_distorted.coil_maps[:, 0].abs().numpy(), - ) - imwrite( - os.path.join( - result_path, - f"brain_sample_{f_idx}_{map_estimation}_{distortion_name}_{reconstructor_name}_reconstructed.tiff", - ), - x_distorted_cropped[0, 0].abs().numpy(), - ) - imwrite( - os.path.join( - result_path, - f"brain_sample_{f_idx}_{map_estimation}_{distortion_name}_{reconstructor_name}_reconstructed.tiff", - ), - x_distorted_cropped[0, 0].abs().numpy(), - ) - except Exception as e: - print(f"Reconstruction with {reconstructor_name} failed due to error: {e}") - continue - - end_recon = datetime.now() - if reconstructor_name not in reconstruction_times: - reconstruction_times[reconstructor_name] = [ - (end_recon - start_recon).total_seconds() - ] - else: - reconstruction_times[reconstructor_name].append( - (end_recon - start_recon).total_seconds() - ) - - end = datetime.now() - if distortion_name not in distortion_times: - distortion_times[distortion_name] = [(end - start).total_seconds()] - else: - distortion_times[distortion_name].append((end - start).total_seconds()) - - -distortion_times_df = pd.DataFrame(distortion_times, index=[sensitivity_map_estimation_algorithm]) -distortion_times_df.to_csv(os.path.join(result_path, "distortion_times.csv"), index=False) -reconstruction_times_df = pd.DataFrame(reconstruction_times) -reconstruction_times_df.to_csv(os.path.join(result_path, "reconstruction_times.csv"), index=False) - -print(distortion_times_df.mean()) -print(reconstruction_times_df.mean()) diff --git a/mri_recon/examples/run_experiments_fastmri_prostateT2.py b/mri_recon/examples/run_experiments_fastmri_prostateT2.py deleted file mode 100644 index e999d56..0000000 --- a/mri_recon/examples/run_experiments_fastmri_prostateT2.py +++ /dev/null @@ -1,289 +0,0 @@ -import os -import sys -import h5py -from tifffile import imwrite -import torch -import deepinv as dinv - -import numpy as np - - -sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) - -from mri_recon.distortions import DistortedKspaceMultiCoilMRI, BaseDistortion -from mri_recon.reconstruction import choose_reconstructor -from mri_recon.utils import Grappa - - -sensitivity_map_estimation_algorithm = [ - # "espirit", - "unity", - # "birdcage", -] - - -DISTORTIONS = [ - # "PhaseEncodeGhosting", - # "CartesianUndersamplingVariableDensityRandom", - "CartesianUndersamplingUniformRandom", - # "HannTaperLP", - # "KaiserTaperLP", - # "RadialHighPassEmphasis", - # "IsotropicLP", - # "OffCenterAnisotropicGaussianBiasField", - # "TranslationMotion", - # "RotationalMotion", - # "SegmentedRotationalMotion", - # "SegmentedTranslationMotion", - "GaussianKspaceBiasField", - # "GaussianNoise", -] - -RECONSTRUCTORS = [ - "zero-filled", - "conjugate-gradient", - # "ram", - # "dip", - # "tv-pgd", - # "wavelet-fista", - # "tv-fista", - # "tv-pdhg", - # "unet", # will trigger download of pretrained weights if not already present - # *list(EXPLICIT_UNET_ALGORITHMS) - #'unet-fastmri', - #'unet-oasis-acceleration4', - #'unet-oasis-acceleration8', - #'unet-oasis-acceleration10', -] - -reference_reconstructor = "conjugate-gradient" -reference_map_estimation = "unity" - -filenames = [ - "/home/melanie.dohmen/mri_recon/data/fastmri/fastMRI_prostate_T2_IDS_001_020/file_prostate_AXT2_001.h5" -] - - -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -result_path = "/home/melanie.dohmen/mri_recon/reports/experiments_fastmri_prostateT2/" - -os.makedirs(result_path, exist_ok=True) - - -def correct_kspace_data_with_calibration( - y: torch.Tensor, calibration_data: torch.Tensor -) -> torch.Tensor: - n_avg, n_slices, n_coils, ksp_w, ksp_h = kspace_data.shape - - # Calib_data shape: num_slices, num_coils, num_pe_cal - grappa_weight_dict = {} - grappa_weight_dict_2 = {} - - kspace_slice_regridded = kspace_data[0, 0, ...] - grappa_obj = Grappa( - np.transpose(kspace_slice_regridded, (2, 0, 1)), kernel_size=(5, 5), coil_axis=1 - ) - - kspace_slice_regridded_2 = kspace_data[1, 0, ...] - grappa_obj_2 = Grappa( - np.transpose(kspace_slice_regridded_2, (2, 0, 1)), kernel_size=(5, 5), coil_axis=1 - ) - - # calculate GRAPPA weights - for slice_num in range(n_slices): - calibration_regridded = calibration_data[slice_num, ...] - grappa_weight_dict[slice_num] = grappa_obj.compute_weights( - np.transpose(calibration_regridded, (2, 0, 1)) - ) - grappa_weight_dict_2[slice_num] = grappa_obj_2.compute_weights( - np.transpose(calibration_regridded, (2, 0, 1)) - ) - - # apply GRAPPA weights - kspace_post_grappa_all = np.zeros(shape=kspace_data.shape, dtype=complex) - - for average, grappa_obj, grappa_weight_dict in zip( - [0, 1, 2], - [grappa_obj, grappa_obj_2, grappa_obj], - [grappa_weight_dict, grappa_weight_dict_2, grappa_weight_dict], - ): - for slice_num in range(n_slices): - kspace_slice_regridded = kspace_data[average, slice_num, ...] - kspace_post_grappa = grappa_obj.apply_weights( - np.transpose(kspace_slice_regridded, (2, 0, 1)), grappa_weight_dict[slice_num] - ) - kspace_post_grappa_all[average, slice_num, ...] = np.moveaxis( - np.moveaxis(kspace_post_grappa, 0, 1), 1, 2 - ) - - return kspace_post_grappa_all - - -for f_idx, filename in enumerate(filenames): - with h5py.File(filename, "r") as hf: - kspace_data = hf["kspace"][:] - calibration_data = hf["calibration_data"][:] - hdr = hf["ismrmrd_header"][()] - reconstruction_rss = hf["reconstruction_rss"][:] - atts = dict() - atts["max"] = hf.attrs["max"] - atts["norm"] = hf.attrs["norm"] - atts["patient_id"] = hf.attrs["patient_id"] - atts["acquisition"] = hf.attrs["acquisition"] - - n_avg, n_slices, n_coils, ksp_w, ksp_h = kspace_data.shape - - # correct k-space data with calibration data: - - kspace_data = correct_kspace_data_with_calibration(kspace_data, calibration_data) - - x = torch.from_numpy(reconstruction_rss).unsqueeze(0).unsqueeze(0) - y = torch.view_as_real(torch.from_numpy(kspace_data)).unsqueeze(0).moveaxis(-1, 1) - # image shape: (1, slices, H, W) - print("image x.shape:", x.shape) - # k-space shape: (1, slices, coils, H, W) - print("k-space y.shape:", y.shape) - x = x.to(device) - y = y.to(device) - - # select middle slice: - x = x[:, :, x.shape[2] // 2, ...] - y = y[:, :, y.shape[2] // 2, ...] - - recon_1 = [] - - for ave in range(n_avg): - # estimate coil maps: - for map_estimation in sensitivity_map_estimation_algorithm: - if map_estimation == "espirit": - estimate_coil_maps = dinv.datasets.MRISliceTransform( - estimate_coil_maps=True, - acs=15, # Num. low frequency, fix to 15 - ) - _, _, params = estimate_coil_maps(target=x[0], kspace=y[0]) - - coil_maps = params["coil_maps"] - print("estimated coil maps shape: ", coil_maps.shape) - imwrite( - os.path.join( - result_path, f"prostate_sample_{f_idx}_{map_estimation}_maps.tiff" - ), - coil_maps[:, 0].abs().numpy(), - ) - - elif map_estimation == "unity": - coil_maps = torch.ones( - (n_coils, ksp_w, ksp_h), dtype=torch.complex64, device=device - ) - print("estimated coil maps shape: ", coil_maps.shape) - imwrite( - os.path.join( - result_path, f"prostate_sample_{f_idx}_{map_estimation}_maps.tiff" - ), - coil_maps[:, 0].abs().numpy(), - ) - - elif map_estimation == "birdcage": - coil_maps = n_coils - - if map_estimation == reference_map_estimation: - physics_clean = dinv.physics.MultiCoilMRI( - img_size=(ksp_w, ksp_h), - mask=None, - coil_maps=coil_maps, - device=device, - ) - - y_distorted = BaseDistortion()(y) - print("base distortion does not change y:", torch.all(y_distorted == y)) - print("kspace distorted.shape: ", y_distorted.shape) - - x_recon_reference = choose_reconstructor(reference_reconstructor)( - y_distorted, physics_clean - ) - print("reconstructed image shape: ", x_recon_reference.shape) - - x_recon_reference_cropped = physics_clean.crop(x_recon_reference, shape=x.shape) - print("cropped reconstructed image shape: ", x_recon_reference_cropped.shape) - - recon_1.append(x_recon_reference_cropped[0, 0].abs().numpy()) - - x_as_inversed_y = physics_clean.A_adjoint(y_distorted, rss=True) - print("inversed y shape: ", x_as_inversed_y.shape) - x_as_inversed_y_cropped = physics_clean.crop(x_as_inversed_y, shape=x.shape) - print("cropped inversed y shape: ", x_as_inversed_y_cropped.shape) - - imwrite( - os.path.join( - result_path, - f"prostate_sample_{f_idx}_{map_estimation}_reconstructed_reference.tiff", - ), - x_recon_reference_cropped[0, 0].abs().numpy(), - ) - imwrite( - os.path.join( - result_path, f"prostate_sample_{f_idx}_{map_estimation}_inversed_y.tiff" - ), - x_as_inversed_y_cropped[0, 0].abs().numpy(), - ) - - if map_estimation == "birdcage": - coil_maps = physics_clean.coil_maps - print("estimated coil maps shape: ", coil_maps.shape) - imwrite( - os.path.join( - result_path, f"prostate_sample_{f_idx}_{map_estimation}_maps.tiff" - ), - coil_maps[:, 0].abs().numpy(), - ) - - for distortion_name in DISTORTIONS: - physics_distorted = DistortedKspaceMultiCoilMRI( - distortion=BaseDistortion(), - img_size=(ksp_w, ksp_h), - mask=None, - coil_maps=coil_maps, - device=device, - ) - - x_distorted_as_inversed_y = physics_distorted.A_adjoint(y_distorted, rss=True) - x_distorted_as_inversed_y_cropped = physics_distorted.crop( - x_distorted_as_inversed_y, shape=x.shape - ) - imwrite( - os.path.join( - result_path, f"prostate_sample_{f_idx}_{distortion_name}_inversed_y.tiff" - ), - x_distorted_as_inversed_y_cropped[0, 0].abs().numpy(), - ) - - for reconstructor_name in RECONSTRUCTORS: - reconstructor = choose_reconstructor(reconstructor_name) - - x_distorted = reconstructor(y_distorted, physics_distorted) - - x_distorted_cropped = physics_distorted.crop(x_distorted, shape=x.shape) - - imwrite( - os.path.join( - result_path, - f"prostate_sample_{f_idx}_{distortion_name}_{reconstructor_name}.tiff", - ), - physics_distorted.coil_maps[:, 0].abs().numpy(), - ) - imwrite( - os.path.join( - result_path, - f"prostate_sample_{f_idx}_{distortion_name}_{reconstructor_name}_reconstructed.tiff", - ), - x_distorted_cropped[0, 0].abs().numpy(), - ) - - recon_1_mean = np.mean(recon_1, axis=0) - imwrite( - os.path.join( - result_path, - f"prostate_sample_{f_idx}_{reference_map_estimation}_reconstructed_reference_mean.tiff", - ), - recon_1_mean, - ) diff --git a/mri_recon/reconstruction/inference.py b/mri_recon/reconstruction/inference.py index 6b9216f..b2c6261 100644 --- a/mri_recon/reconstruction/inference.py +++ b/mri_recon/reconstruction/inference.py @@ -48,7 +48,7 @@ def compatible_dataset_with_reconstructor(dataset: str, reconstructor_name: str) # fast mri u-net is only trained with knee data if reconstructor_name == FASTMRI_UNET_ALGORITHM: - if (dataset == "fastmri_knee"): + if dataset == "fastmri_knee": return True else: return False @@ -59,7 +59,7 @@ def compatible_dataset_with_reconstructor(dataset: str, reconstructor_name: str) return True else: return False - + # all other (classic) reconstructors work with any dataset: else: return True @@ -91,7 +91,9 @@ def choose_reconstructor( """ if dataset is not None and not compatible_dataset_with_reconstructor(dataset, name): - raise ValueError(f"Reconstructor {name} is not compatible with dataset {dataset}, because it was trained with a different image domain.") + raise ValueError( + f"Reconstructor {name} is not compatible with dataset {dataset}, because it was trained with a different image domain." + ) match name: case "zero-filled": diff --git a/mri_recon/utils/__init__.py b/mri_recon/utils/__init__.py index 0436aa5..a0648ad 100644 --- a/mri_recon/utils/__init__.py +++ b/mri_recon/utils/__init__.py @@ -15,7 +15,6 @@ from .oasis_adapter import image_to_kspace as image_to_kspace from .oasis_adapter import kspace_to_image as kspace_to_image from .oasis_adapter import image_to_fastmri_measurement as image_to_fastmri_measurement -from .oasis_adapter import oasis_kspace_to_fastmri_measurement as oasis_kspace_to_fastmri_measurement from .prostate_adaptor import FastMRIProstateDataset as FastMRIProstateDataset from .plot import save_kspace_plot as save_kspace_plot from .plot import _kspace_to_log_magnitude as _kspace_to_log_magnitude diff --git a/mri_recon/utils/oasis_adapter.py b/mri_recon/utils/oasis_adapter.py index f9542b8..88ba5de 100644 --- a/mri_recon/utils/oasis_adapter.py +++ b/mri_recon/utils/oasis_adapter.py @@ -9,7 +9,7 @@ from torch.utils.data import Dataset import deepinv as dinv -from mri_recon.distortions import BaseDistortion, DistortedKspaceMultiCoilMRI +from mri_recon.distortions import BaseDistortion class OasisSliceDataset(Dataset): diff --git a/mri_recon/utils/plot.py b/mri_recon/utils/plot.py index 4f18b1e..35921b7 100644 --- a/mri_recon/utils/plot.py +++ b/mri_recon/utils/plot.py @@ -77,7 +77,7 @@ def save_kspace_plot( def convert_image_for_save(im: torch.Tensor) -> np.ndarray: """ - Convert a PyTorch tensor image complex tensor to a real-valued NumPy array + Convert a PyTorch tensor image complex tensor to a real-valued NumPy array by calculating the magnitude. (B, 2, H, W) or (B, H, W) with complex type -> (B, H, W) diff --git a/mri_recon/utils/prostate_adaptor.py b/mri_recon/utils/prostate_adaptor.py index 5145fcd..c7b1d6a 100644 --- a/mri_recon/utils/prostate_adaptor.py +++ b/mri_recon/utils/prostate_adaptor.py @@ -4,8 +4,11 @@ import numpy as np import torch + class FastMRIProstateDataset(torch.utils.data.Dataset): - def __init__(self, data_path: str, num_samples: None | int = None, slice_index: str="middle") -> None: + def __init__( + self, data_path: str, num_samples: None | int = None, slice_index: str = "middle" + ) -> None: self.data_path = data_path self.num_samples = num_samples self.slice_index = slice_index @@ -30,9 +33,11 @@ def get_image_data(self) -> np.ndarray: continue if self.slice_index == "middle": - image_result_list.append(image_recon[image_recon.shape[0]//2]) + image_result_list.append(image_recon[image_recon.shape[0] // 2]) else: - image_result_list.extend([image_recon[i, :, :] for i in range(image_recon.shape[0])]) + image_result_list.extend( + [image_recon[i, :, :] for i in range(image_recon.shape[0])] + ) return image_result_list From b4d21a438d94714e743252aca59e6ea468cdbd54 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Mon, 1 Jun 2026 10:29:33 +0000 Subject: [PATCH 06/22] remove personal config file --- examples/config_md.yaml | 45 ----------------------------------------- 1 file changed, 45 deletions(-) delete mode 100644 examples/config_md.yaml diff --git a/examples/config_md.yaml b/examples/config_md.yaml deleted file mode 100644 index a85bc17..0000000 --- a/examples/config_md.yaml +++ /dev/null @@ -1,45 +0,0 @@ -data: - #"fastmri_knee": "/home/melanie.dohmen/ArtifactLab/data/singlecoil_val" - "oasis": "/home/melanie.dohmen/ArtifactLab/data/oasis" - "fastmri_brain": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_multicoil_brain_test" - "cmrxrecon": "/home/melanie.dohmen/ArtifactLab/data/CMRxRecon" - "fastmri_prostate": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_prostate_T2_IDS_001_020" -distortions: - - "BaseDistortion" - #- "CartesianUndersamplingVariableDensity" - #- "CartesianUndersamplingUniformRandom" - #- "CartesianUndersamplingUniformRandomZeroACS" - #- "CartesianUndersamplingEquispaced" - #- "CartesianUndersamplingEquispacedZeroACS" - #- "PartialFourier" - #- "PhaseEncodeGhosting" - #- "SegmentedTranslationMotion" - #- "SegmentedRotationalMotion" - #- "TranslationMotion" - #- "RotationalMotion" - #- "OffCenterAnisotropicGaussianBiasField" - #- "GaussianBiasField" - #- "AnisotropicLP" - #- "HannTaperLP" - #- "KaiserTaperLP" - #- "GaussianNoise" - #- "IsotropicLP" - #- "RadialHighPassEmphasis" -reconstruction_algorithms: - - "zero-filled" - - "conjugate-gradient" - - "ram" - - "dip" - - "tv-pgd" - - "wavelet-fista" - - "tv-fista" - - "tv-pdhg" - - "unet-fastmri" - - "unet-oasis-acceleration4" - - "unet-oasis-acceleration8" - - "unet-oasis-acceleration10" -num_samples: 1 -keep_fraction: 0.25 -center_fraction: 0.125 -verbose: true -results_dir: "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1" \ No newline at end of file From 501e522fa3aa530ed4b5d3d0bf6d3e3a0a7f4c28 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Mon, 1 Jun 2026 10:35:28 +0000 Subject: [PATCH 07/22] add test to transform images to fastmri like k-space data --- tests/test_utils_io.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_utils_io.py b/tests/test_utils_io.py index 304446d..84b0e11 100644 --- a/tests/test_utils_io.py +++ b/tests/test_utils_io.py @@ -8,6 +8,8 @@ fastmri_measurement_to_image, fastmri_measurement_to_oasis_kspace, kspace_to_image, + oasis_kspace_to_fastmri_measurement, + image_to_fastmri_measurement, ) from mri_recon.utils.io import download_file_with_sha256, download_google_drive_file_with_sha256 @@ -99,3 +101,19 @@ def test_fastmri_measurement_helpers_match_centered_oasis_path(): atol=1e-6, rtol=1e-6, ) + + y_fastmri_from_oasis = oasis_kspace_to_fastmri_measurement(y_oasis, device="cpu") + assert torch.allclose( + y_fastmri_from_oasis, + y_fastmri, + atol=1e-6, + rtol=1e-6, + ) + + y_fastmri_from_image = image_to_fastmri_measurement(x_native, device="cpu") + assert torch.allclose( + y_fastmri_from_image, + y_fastmri, + atol=1e-6, + rtol=1e-6, + ) From 07d0da814776e7d879f42fbdc3f7a67fffdfc584 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Mon, 1 Jun 2026 10:42:04 +0000 Subject: [PATCH 08/22] correct times for reconstructions, remove debugging print statements --- examples/run_all.py | 31 ++++++++----------------------- 1 file changed, 8 insertions(+), 23 deletions(-) diff --git a/examples/run_all.py b/examples/run_all.py index 0141a68..80f4f34 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -85,11 +85,8 @@ def get_measurement_sample( elif dataset_name == "cmrxrecon": # reference image, shape: (B, 2, n_timepoints, (n_coils), H, W) x = sample_batch[0].to(run_device) - print(f"\t[Debug] Reference image shape: {x.shape}, dtype: {x.dtype}") # k-space data, shape: (B, 2, n_timepoints, (n_coils), H, W) dtype: float32 y = sample_batch[1].to(run_device) - print(f"\t[Debug] k-space shape: {y.shape}, dtype: {y.dtype}") - # not available for all samples, either None or # shape (1, num_coils, H, W) coil_maps = ( @@ -99,36 +96,22 @@ def get_measurement_sample( and "coil_maps" in sample_batch[2] else None ) - print( - f"\t[Debug] Coil maps shape: {coil_maps.shape if coil_maps is not None else None}, dtype: {coil_maps.dtype if coil_maps is not None else None}" - ) - # centered k-space data, shape: (B, 2, H, W) dtype: float32 + + # centered k-space data, shape: (B, 2, num_clois, H, W) dtype: float32 y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) - print(f"\t[Debug] Centered k-space shape: {y_centered.shape}, dtype: {y_centered.dtype}") - # reconstruct coil-combined image reference from multi-coil k-space data using - # integrated espirit sensitivity map estimation, RSS coil combination - # x = fastmri_measurement_to_image(y, coil_maps=coil_maps, rss=True) - # print(f"\t[Debug] Reference image shape: {x.shape}, dtype: {x.dtype}") elif dataset_name == "fastmri_prostate": - # reference image, shape: (slices, W, H): dtype float32 + # reference image, shape: (B, W, H): dtype float32 x = sample_batch[0].to(run_device) - print(f"\t[Debug] Reference image shape: {x.shape}, type: {x.dtype}") # add zero imaginary channel: - # (B, slices, H, W) -> (B, 2, slices, H, W) + # (B, H, W) -> (B, 2, H, W) x = torch.stack([x, torch.zeros_like(x)], dim=1) - # y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) - # create oasis-like k-space data from image: + # (B, 2, H, W) y_centered = image_to_kspace(x) - print(f"\t[Debug] Centered k-space shape: {y_centered.shape}, type: {y_centered.dtype}") y = oasis_kspace_to_fastmri_measurement(y_centered) - print(f"\tk-space shape {y.shape}[{y.dtype}] and reference image shape: {x.shape}[{x.dtype}]") - if coil_maps is not None: - print(f"\tcoil maps shape: {coil_maps.shape}[{coil_maps.dtype}]") - return x, y, y_centered, coil_maps @@ -284,7 +267,7 @@ def get_measurement_sample( ), convert_image_for_save(x_corrected), ) - print(f"\t\t... done in {start - datetime.now()}") + print(f"\t\t... done in {datetime.now() - start}") except Exception as e: print(f"Error using {reconstructor_name}: {e}") @@ -314,6 +297,7 @@ def get_measurement_sample( # skip all reconstructors, that don't use the oasis-centered path if uses_oasis_centered_path(dataset_name, reconstructor_name): print(f"\t\t{reconstructor_name} ...") + start = datetime.now() if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): reconstructor = choose_reconstructor( reconstructor_name, @@ -368,6 +352,7 @@ def get_measurement_sample( ), convert_image_for_save(x_corrected), ) + print(f"\t\t... done in {datetime.now() - start}") except Exception as e: print( From 0d344882f87284a548cad831d58695c2ce63d884 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Mon, 1 Jun 2026 10:59:35 +0000 Subject: [PATCH 09/22] correct type of dataset - algorithm validation, centered FFT path for oasis model only --- examples/fastmri_inference_plot.py | 11 ++++++++--- examples/run_all.py | 4 ++-- mri_recon/reconstruction/inference.py | 7 +------ tests/test_reconstructions.py | 23 ++++++++++++----------- 4 files changed, 23 insertions(+), 22 deletions(-) diff --git a/examples/fastmri_inference_plot.py b/examples/fastmri_inference_plot.py index b7a1913..d580324 100644 --- a/examples/fastmri_inference_plot.py +++ b/examples/fastmri_inference_plot.py @@ -38,7 +38,7 @@ OASISSinglecoilUnetReconstructor, choose_reconstructor, uses_oasis_centered_path, - validate_algorithm_dataset_compatibility, + compatible_dataset_with_reconstructor, EXPLICIT_UNET_ALGORITHMS, ) from mri_recon.utils import ( @@ -376,8 +376,13 @@ def build_physics_pair( selected_algorithms = ALGORITHMS if args.algorithm == "" else [args.algorithm] selected_distortions = DISTORTIONS if args.distortion == "" else [args.distortion] - for algo_name in selected_algorithms: - validate_algorithm_dataset_compatibility(args.dataset, algo_name) + + # skip non-compatible algorithm-dataset pairs + selected_algorithms = [ + algo_name + for algo_name in selected_algorithms + if compatible_dataset_with_reconstructor(args.dataset, algo_name) + ] # set up report dir if args.dataset == "fastmri": diff --git a/examples/run_all.py b/examples/run_all.py index 80f4f34..e051129 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -209,7 +209,7 @@ def get_measurement_sample( for reconstructor_name in config["reconstruction_algorithms"]: # only run on reconstructors, that use the fastmri-like k-space - if not uses_oasis_centered_path(dataset_name, reconstructor_name): + if not uses_oasis_centered_path(reconstructor_name): print(f"\t\t{reconstructor_name} ...") start = datetime.now() if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): @@ -295,7 +295,7 @@ def get_measurement_sample( for reconstructor_name in config["reconstruction_algorithms"]: # skip all reconstructors, that don't use the oasis-centered path - if uses_oasis_centered_path(dataset_name, reconstructor_name): + if uses_oasis_centered_path(reconstructor_name): print(f"\t\t{reconstructor_name} ...") start = datetime.now() if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): diff --git a/mri_recon/reconstruction/inference.py b/mri_recon/reconstruction/inference.py index b2c6261..8155ff6 100644 --- a/mri_recon/reconstruction/inference.py +++ b/mri_recon/reconstruction/inference.py @@ -28,18 +28,13 @@ def uses_oasis_centered_path( - dataset: str, algorithm: str, ) -> bool: """Return whether inference should use the centered OASIS k-space path. - OASIS samples always use the centered FFT convention. FastMRI only switches - to that path when the selected algorithm is one of the explicit OASIS U-Net + The centered FFT path is only used with the explicit OASIS U-Net variants. """ - - if dataset == "oasis": - return True return algorithm in OASIS_UNET_ALGORITHMS diff --git a/tests/test_reconstructions.py b/tests/test_reconstructions.py index 3c95455..9ca3747 100644 --- a/tests/test_reconstructions.py +++ b/tests/test_reconstructions.py @@ -17,7 +17,7 @@ OASIS_UNET_ALGORITHMS, choose_reconstructor, uses_oasis_centered_path, - validate_algorithm_dataset_compatibility, + compatible_dataset_with_reconstructor, ) from mri_recon.distortions import DistortedKspaceMultiCoilMRI @@ -245,21 +245,22 @@ def fake_resolve(_cls, acceleration, manifest_path=None): def test_validate_algorithm_dataset_compatibility_accepts_supported_explicit_unets(): - validate_algorithm_dataset_compatibility("fastmri", FASTMRI_UNET_ALGORITHM) - validate_algorithm_dataset_compatibility("fastmri", "unet-oasis-acceleration8") - validate_algorithm_dataset_compatibility("oasis", "unet-oasis-acceleration4") + assert compatible_dataset_with_reconstructor("fastmri", FASTMRI_UNET_ALGORITHM) + assert compatible_dataset_with_reconstructor("fastmri", "unet-oasis-acceleration8") + assert compatible_dataset_with_reconstructor("oasis", "unet-oasis-acceleration4") def test_validate_algorithm_dataset_compatibility_rejects_unsupported_oasis_fastmri_combo(): - with pytest.raises(ValueError, match="unet-fastmri"): - validate_algorithm_dataset_compatibility("oasis", FASTMRI_UNET_ALGORITHM) + assert not compatible_dataset_with_reconstructor("oasis", FASTMRI_UNET_ALGORITHM) + assert not compatible_dataset_with_reconstructor("oasis", FASTMRI_UNET_ALGORITHM) + assert not compatible_dataset_with_reconstructor("oasis", FASTMRI_UNET_ALGORITHM) + assert not compatible_dataset_with_reconstructor("oasis", FASTMRI_UNET_ALGORITHM) -def test_uses_oasis_centered_path_tracks_dataset_and_explicit_algorithm(): - assert uses_oasis_centered_path("oasis", FASTMRI_UNET_ALGORITHM) is True - assert uses_oasis_centered_path("fastmri", "unet-oasis-acceleration8") is True - assert uses_oasis_centered_path("fastmri", FASTMRI_UNET_ALGORITHM) is False - assert uses_oasis_centered_path("fastmri", "tv-pgd") is False +def test_uses_oasis_centered_algorithm(): + assert uses_oasis_centered_path(FASTMRI_UNET_ALGORITHM) is False + assert uses_oasis_centered_path("unet-oasis-acceleration4") is True + assert uses_oasis_centered_path("tv-pgd") is False def test_choose_reconstructor_selects_oasis_unet_for_fastmri_when_requested(monkeypatch): From 296dd8eac0c62217c7956177c75fa6666e792aae Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Mon, 1 Jun 2026 11:22:24 +0000 Subject: [PATCH 10/22] update tests for validating data + reconstruction pairs, update test for transforming oasis -lik data to fast-MRI like k-space data --- tests/test_reconstructions.py | 16 ++++++++-------- tests/test_utils_io.py | 9 ++++----- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/tests/test_reconstructions.py b/tests/test_reconstructions.py index 9ca3747..966ad49 100644 --- a/tests/test_reconstructions.py +++ b/tests/test_reconstructions.py @@ -245,16 +245,16 @@ def fake_resolve(_cls, acceleration, manifest_path=None): def test_validate_algorithm_dataset_compatibility_accepts_supported_explicit_unets(): - assert compatible_dataset_with_reconstructor("fastmri", FASTMRI_UNET_ALGORITHM) - assert compatible_dataset_with_reconstructor("fastmri", "unet-oasis-acceleration8") + assert compatible_dataset_with_reconstructor("fastmri_knee", FASTMRI_UNET_ALGORITHM) + assert compatible_dataset_with_reconstructor("fastmri_brain", "unet-oasis-acceleration8") assert compatible_dataset_with_reconstructor("oasis", "unet-oasis-acceleration4") def test_validate_algorithm_dataset_compatibility_rejects_unsupported_oasis_fastmri_combo(): assert not compatible_dataset_with_reconstructor("oasis", FASTMRI_UNET_ALGORITHM) - assert not compatible_dataset_with_reconstructor("oasis", FASTMRI_UNET_ALGORITHM) - assert not compatible_dataset_with_reconstructor("oasis", FASTMRI_UNET_ALGORITHM) - assert not compatible_dataset_with_reconstructor("oasis", FASTMRI_UNET_ALGORITHM) + assert not compatible_dataset_with_reconstructor("fastmri_brain", FASTMRI_UNET_ALGORITHM) + assert not compatible_dataset_with_reconstructor("fastmri_prostate", FASTMRI_UNET_ALGORITHM) + assert not compatible_dataset_with_reconstructor("fastmri_knee", "unet-oasis-acceleration10") def test_uses_oasis_centered_algorithm(): @@ -281,7 +281,7 @@ def fake_oasis(*, acceleration, device): reconstructor = choose_reconstructor( "unet-oasis-acceleration8", - dataset="fastmri", + dataset="oasis", device="cpu", ) @@ -303,7 +303,7 @@ def fake_fastmri(*, device): reconstructor = choose_reconstructor( FASTMRI_UNET_ALGORITHM, - dataset="fastmri", + dataset="fastmri_knee", device="cpu", ) @@ -328,7 +328,7 @@ def fake_oasis(*, acceleration, device): for algorithm_name in OASIS_UNET_ALGORITHMS: reconstructor = choose_reconstructor( algorithm_name, - dataset="fastmri", + dataset="fastmri_brain", device="cpu", ) assert isinstance(reconstructor, Marker) diff --git a/tests/test_utils_io.py b/tests/test_utils_io.py index 84b0e11..9192da1 100644 --- a/tests/test_utils_io.py +++ b/tests/test_utils_io.py @@ -2,8 +2,8 @@ from io import BytesIO import torch +import deepinv as dinv -from mri_recon.distortions import BaseDistortion, DistortedKspaceMultiCoilMRI from mri_recon.utils.oasis_adapter import ( fastmri_measurement_to_image, fastmri_measurement_to_oasis_kspace, @@ -80,8 +80,7 @@ def fake_urlopen(url, timeout=30): def test_fastmri_measurement_helpers_match_centered_oasis_path(): x = torch.randn(1, 2, 16, 12) - physics = DistortedKspaceMultiCoilMRI( - distortion=BaseDistortion(), + physics = dinv.physics.MultiCoilMRI( img_size=(1, 2, *x.shape[-2:]), device="cpu", ) @@ -102,7 +101,7 @@ def test_fastmri_measurement_helpers_match_centered_oasis_path(): rtol=1e-6, ) - y_fastmri_from_oasis = oasis_kspace_to_fastmri_measurement(y_oasis, device="cpu") + y_fastmri_from_oasis = oasis_kspace_to_fastmri_measurement(y_oasis) assert torch.allclose( y_fastmri_from_oasis, y_fastmri, @@ -110,7 +109,7 @@ def test_fastmri_measurement_helpers_match_centered_oasis_path(): rtol=1e-6, ) - y_fastmri_from_image = image_to_fastmri_measurement(x_native, device="cpu") + y_fastmri_from_image = image_to_fastmri_measurement(x_native) assert torch.allclose( y_fastmri_from_image, y_fastmri, From 98c13636f2d5c22a2effcf97c25a187dd2fcbca8 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Thu, 4 Jun 2026 11:32:03 +0000 Subject: [PATCH 11/22] correct distortion names, allow num_samples = None, remove commented out code --- examples/config.yaml | 4 ++-- examples/run_all.py | 15 ++++++++++----- mri_recon/distortions/base.py | 3 --- mri_recon/utils/oasis_adapter.py | 11 ----------- 4 files changed, 12 insertions(+), 21 deletions(-) diff --git a/examples/config.yaml b/examples/config.yaml index 7fbed4e..6ab8c1c 100644 --- a/examples/config.yaml +++ b/examples/config.yaml @@ -18,8 +18,8 @@ distortions: - "SegmentedRotationalMotion" - "TranslationMotion" - "RotationalMotion" - - "OffCenterAnisotropicGaussianBiasField" - - "GaussianBiasField" + - "OffCenterAnisotropicGaussianKspaceBiasField" + - "GaussianKspaceBiasField" - "AnisotropicLP" - "HannTaperLP" - "KaiserTaperLP" diff --git a/examples/run_all.py b/examples/run_all.py index e051129..abebc8c 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -87,8 +87,9 @@ def get_measurement_sample( x = sample_batch[0].to(run_device) # k-space data, shape: (B, 2, n_timepoints, (n_coils), H, W) dtype: float32 y = sample_batch[1].to(run_device) - # not available for all samples, either None or - # shape (1, num_coils, H, W) + + # maybe not needed, as there are no coil maps in the current + # cmrxrecon sample coil_maps = ( sample_batch[2]["coil_maps"].to(run_device) if isinstance(sample_batch, (tuple, list)) @@ -96,10 +97,14 @@ def get_measurement_sample( and "coil_maps" in sample_batch[2] else None ) - - # centered k-space data, shape: (B, 2, num_clois, H, W) dtype: float32 + # centered k-space data, shape: (B, 2, n_timepoints, (n_coils), H, W) dtype: float32 y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) + # select the first timepoint in order to simplify the evaluation of the reconstruction algorithms + y = y[:, :, 0, ...] + y_centered = y_centered[:, :, 0, ...] + x = x[:, :, 0, ...] + elif dataset_name == "fastmri_prostate": # reference image, shape: (B, W, H): dtype float32 x = sample_batch[0].to(run_device) @@ -164,7 +169,7 @@ def get_measurement_sample( # loop through samples of dataset for i, batch in enumerate(iter(torch.utils.data.DataLoader(dataset))): # exit loop if we have processed the specified number of samples - if i >= config["num_samples"]: + if (config["num_samples"] is not None) and (i >= config["num_samples"]): break print(f"{dataset_name} sample {i}...") diff --git a/mri_recon/distortions/base.py b/mri_recon/distortions/base.py index d89efd0..c5b0734 100644 --- a/mri_recon/distortions/base.py +++ b/mri_recon/distortions/base.py @@ -188,6 +188,3 @@ def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: y = self.distortion.A_adjoint(y) return super().A_adjoint(y) - - # def A_dagger(self, y: torch.Tensor, **kwargs) -> torch.Tensor: - # return super().A_dagger(y, coil_maps=self.coil_maps, **kwargs) diff --git a/mri_recon/utils/oasis_adapter.py b/mri_recon/utils/oasis_adapter.py index 88ba5de..7903805 100644 --- a/mri_recon/utils/oasis_adapter.py +++ b/mri_recon/utils/oasis_adapter.py @@ -392,14 +392,3 @@ def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: """ return kspace_to_image(self.distortion.A_adjoint(y)) - - # def A_dagger(self, y: torch.Tensor, **kwargs) -> torch.Tensor: - # r""" - # Computes least squares solution to the MRI inverse problem, as proposed in `SENSE: Sensitivity encoding for fast MRI `_. - - # By default uses conjugate gradient solver. Overwrite default solver arguments by passing `kwargs`. See :func:`deepinv.optim.linear.least_squares` for details. - - # :param dict kwargs: kwargs to pass to base :meth:`deepinv.physics.LinearPhysics.A_dagger`. - # :returns: (:class:`torch.Tensor`) image with shape `(B,2,...,H,W)` - # """ - # return super().A_dagger(y, **kwargs) From 0e7e40a7b6fd8f81db94cf836d06eae4c42f66d7 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Mon, 8 Jun 2026 15:02:12 +0000 Subject: [PATCH 12/22] add params to distortions, add Biasfield in image domain, add N4 bias Field correction --- examples/config.yaml | 81 ++++-- examples/fastmri_inference_plot.py | 386 ++++++++++++----------------- examples/run_all.py | 379 ++++++++++++++++------------ mri_recon/distortions/__init__.py | 8 +- mri_recon/distortions/base.py | 44 ++++ mri_recon/distortions/biasfield.py | 132 +++++++++- mri_recon/distortions/ghosting.py | 8 +- mri_recon/distortions/motion.py | 30 ++- mri_recon/distortions/utils.py | 89 ++++++- pyproject.toml | 1 + uv.lock | 20 ++ 11 files changed, 763 insertions(+), 415 deletions(-) diff --git a/examples/config.yaml b/examples/config.yaml index 6ab8c1c..65f6d84 100644 --- a/examples/config.yaml +++ b/examples/config.yaml @@ -6,26 +6,64 @@ data: "fastmri_prostate": "/path/to/fastmri/fastMRI_prostate_T2_IDS_001_020" distortions: - - "BaseDistortion" - - "CartesianUndersamplingVariableDensity" - - "CartesianUndersamplingUniformRandom" - - "CartesianUndersamplingUniformRandomZeroACS" - - "CartesianUndersamplingEquispaced" - - "CartesianUndersamplingEquispacedZeroACS" - - "PartialFourier" - - "PhaseEncodeGhosting" - - "SegmentedTranslationMotion" - - "SegmentedRotationalMotion" - - "TranslationMotion" - - "RotationalMotion" - - "OffCenterAnisotropicGaussianKspaceBiasField" - - "GaussianKspaceBiasField" - - "AnisotropicLP" - - "HannTaperLP" - - "KaiserTaperLP" - - "GaussianNoise" - - "IsotropicLP" - - "RadialHighPassEmphasis" + - "BaseDistortion": {} + - "CartesianUndersamplingVariableDensity": + "keep_fraction": "keep_fraction" + "center_fraction": "center_fraction" + - "CartesianUndersamplingUniformRandom": + "keep_fraction": "keep_fraction" + "center_fraction": "center_fraction" + - "CartesianUndersamplingUniformRandomZeroACS": + "keep_fraction": "keep_fraction" + - "CartesianUndersamplingEquispaced": + "keep_fraction": "keep_fraction" + "center_fraction": "center_fraction" + - "CartesianUndersamplingEquispacedZeroACS": + "keep_fraction": "keep_fraction" + - "PartialFourier": + "side": "high" + - "PhaseEncodeGhosting": + "line_period": 2, + "line_offset": 1, + "phase_error_radians": torch.pi / 2, + "corrupted_line_scale": 1.0 + - "SegmentedTranslationMotion": + "shift_x_pixels": [0.0, 20.0, 50.0, -50.0] + "shift_y_pixels": [0.0, 10.0, -20.0, 20.0] + - "SegmentedRotationalMotion": + #"angle_radians": [0.0, torch.pi / 20, -torch.pi / 24, torch.pi / 16] + "angle_degrees": [0.0, 18.0, -15.0, 22.5] + - "TranslationMotion": + "shift_x_pixels": 60 + "shift_y_pixels": 10 + - "RotationalMotion": + # angle_radians=torch.pi / 6 + "angle_degrees": 60.0 + - "OffCenterAnisotropicGaussianBiasField": + "width_x_fraction": 0.2 + "width_y_fraction": 0.35 + "center_x_fraction": 0.15 + "center_y_fraction": -0.1 + "edge_gain": 0.05 + - "GaussianBiasField": + "width_fraction": 0.35 + "edge_gain": 0.05 + - "AnisotropicLP": + "kx_radius_fraction": 1.0, + "ky_radius_fraction": 0.25, + - "HannTaperLP": + "radius_fraction": 0.35, + "transition_fraction": 0.4, + - "KaiserTaperLP": + "radius_fraction": 0.35, + "transition_fraction": 0.4, + "beta": 8.6, + - "GaussianNoise": + "sigma": 0.00001 + - "IsotropicLP": + "radius_fraction": 0.1 + - "RadialHighPassEmphasis": + "alpha": 0.4 reconstruction_algorithms: - "zero-filled" - "conjugate-gradient" @@ -39,8 +77,7 @@ reconstruction_algorithms: - "unet-oasis-acceleration4" #- "unet-oasis-acceleration8" #- "unet-oasis-acceleration10" +add_N4Correction: false num_samples: 1 -keep_fraction: 0.25 -center_fraction: 0.125 verbose: true results_dir: "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1" diff --git a/examples/fastmri_inference_plot.py b/examples/fastmri_inference_plot.py index d580324..9b34fd8 100644 --- a/examples/fastmri_inference_plot.py +++ b/examples/fastmri_inference_plot.py @@ -15,23 +15,9 @@ import torch from mri_recon.distortions import ( - AnisotropicResolutionReduction, + choose_distortion_with_params, BaseDistortion, - CartesianUndersampling, DistortedKspaceMultiCoilMRI, - GaussianKspaceBiasField, - GaussianNoiseDistortion, - HannTaperResolutionReduction, - IsotropicResolutionReduction, - KaiserTaperResolutionReduction, - OffCenterAnisotropicGaussianKspaceBiasField, - PartialFourierDistortion, - PhaseEncodeGhostingDistortion, - RadialHighPassEmphasisDistortion, - RotationalMotionDistortion, - SegmentedRotationalMotionDistortion, - SegmentedTranslationMotionDistortion, - TranslationMotionDistortion, ) from mri_recon.reconstruction import ( ConjugateGradientReconstructor, @@ -72,27 +58,91 @@ ] DISTORTIONS = [ - "no distortion", - # "Cartesian undersampling (variable density)", - # "Cartesian undersampling (uniform random)", - # "Cartesian undersampling (uniform random, zero ACS)", - # "Cartesian undersampling (equispaced)", - "Cartesian undersampling (equispaced, zero ACS)", - # "Partial Fourier", - # "Phase-encode ghosting", - # "Segmented translation motion", - # "Segmented rotational motion", - # "Translation motion", - # "Rotational motion", - # "Off-center anisotropic Gaussian bias field", - # "Gaussian bias field", - # "Anisotropic LP", - # "Hann taper LP", - # "Kaiser taper LP", - # "Gaussian noise", - # "Isotropic LP", - # "Radial high-pass emphasis", + {"BaseDistortion": {}}, + { + "CartesianUndersamplingVariableDensity": { + "keep_fraction": 0.25, + "center_fraction": 0.125, + } + }, + # {"CartesianUndersamplingUniformRandom": { + # "keep_fraction": 0.25, + # "center_fraction": 0.125, + # }}, + # {"CartesianUndersamplingUniformRandomZeroACS": { + # "keep_fraction": 0.25, + # }}, + # {"CartesianUndersamplingEquispaced": { + # "keep_fraction": 0.25, + # "center_fraction": 0.125, + # }}, + # {"CartesianUndersamplingEquispacedZeroACS": { + # "keep_fraction": 0.25, + # }}, + # {"PartialFourier": { + # "side": "high", + # }}, + # {"PhaseEncodeGhosting": { + # "line_period": 2, + # "line_offset": 1, + # "phase_error_degrees": 90.0, + # "corrupted_line_scale": 1.0, + # }}, + # {"SegmentedTranslationMotion": { + # "shift_x_pixels": [0.0, 20.0, 50.0, -50.0], + # "shift_y_pixels": [0.0, 10.0, -20.0, 20.0], + # }}, + # {"SegmentedRotationalMotion": { + # #"angle_radians": [0.0, torch.pi / 20, -torch.pi / 24, torch.pi / 16] + # "angle_degrees": [0.0, 18.0, -15.0, 22.5], + # }}, + # {"TranslationMotion": { + # "shift_x_pixels": 60, + # "shift_y_pixels": 10, + # }}, + # {"RotationalMotion": { + # # angle_radians=torch.pi / 6 + # "angle_degrees": 60.0, + # }}, + { + "OffCenterAnisotropicGaussianBiasField": { + "width_x_fraction": 0.2, + "width_y_fraction": 0.35, + "center_x_fraction": 0.15, + "center_y_fraction": -0.1, + "edge_gain": 0.3, + } + }, + { + "GaussianBiasField": { + "width_fraction": 0.35, + "edge_gain": 0.4, + } + }, + # {"AnisotropicLP": { + # "kx_radius_fraction": 1.0, + # "ky_radius_fraction": 0.25, + # }}, + # {"HannTaperLP": { + # "radius_fraction": 0.35, + # "transition_fraction": 0.4, + # }}, + # {"KaiserTaperLP": { + # "radius_fraction": 0.35, + # "transition_fraction": 0.4, + # "beta": 8.6, + # }}, + # {"GaussianNoise": { + # "sigma": 0.00001, + # }}, + # {"IsotropicLP": { + # "radius_fraction": 0.1, + # }}, + # {"RadialHighPassEmphasis": { + # "alpha": 0.4, + # }} ] + METRICS = [ "PSNR", # "NMSE", @@ -103,125 +153,6 @@ ] -def choose_distortion( - name: str, - keep_fraction: float = 0.25, - center_fraction: float = 0.125, - cartesian_axis: int = -2, -) -> BaseDistortion: - """Build one distortion operator for the inference comparison script. - - The ``cartesian_axis`` is supplied by the active measurement convention: - FastMRI-native runs use the repository's existing axis, while OASIS-native - and FastMRI-to-OASIS runs use the centered OASIS axis. - """ - - match name: - case "Phase-encode ghosting": - return PhaseEncodeGhostingDistortion( - line_period=2, - line_offset=1, - phase_error_radians=torch.pi / 2, - corrupted_line_scale=1.0, - ) - case "Cartesian undersampling (variable density)": - return CartesianUndersampling( - keep_fraction=keep_fraction, - center_fraction=center_fraction, - pattern="variable_density_random", - axis=cartesian_axis, - seed=42, - ) - case "Cartesian undersampling (uniform random)": - return CartesianUndersampling( - keep_fraction=keep_fraction, - center_fraction=center_fraction, - pattern="uniform_random", - axis=cartesian_axis, - seed=42, - ) - case "Cartesian undersampling (uniform random, zero ACS)": - return CartesianUndersampling( - keep_fraction=keep_fraction, - center_fraction=0.0, - pattern="uniform_random", - axis=cartesian_axis, - seed=42, - ) - case "Cartesian undersampling (equispaced)": - return CartesianUndersampling( - keep_fraction=keep_fraction, - center_fraction=center_fraction, - pattern="equispaced", - axis=cartesian_axis, - seed=42, - ) - case "Cartesian undersampling (equispaced, zero ACS)": - return CartesianUndersampling( - keep_fraction=keep_fraction, - center_fraction=0.0, - pattern="equispaced", - axis=cartesian_axis, - seed=42, - ) - case "Partial Fourier": - return PartialFourierDistortion( - partial_fraction=0.7, - center_fraction=center_fraction, - axis=cartesian_axis, - side="high", - ) - case "Anisotropic LP": - return AnisotropicResolutionReduction( - kx_radius_fraction=1.0, - ky_radius_fraction=0.25, - ) - case "Hann taper LP": - return HannTaperResolutionReduction( - radius_fraction=0.35, - transition_fraction=0.4, - ) - case "Kaiser taper LP": - return KaiserTaperResolutionReduction( - radius_fraction=0.35, - transition_fraction=0.4, - beta=8.6, - ) - case "Radial high-pass emphasis": - return RadialHighPassEmphasisDistortion(alpha=0.4) - case "Isotropic LP": - return IsotropicResolutionReduction(radius_fraction=0.1) - case "Off-center anisotropic Gaussian bias field": - return OffCenterAnisotropicGaussianKspaceBiasField( - width_x_fraction=0.2, - width_y_fraction=0.35, - center_x_fraction=0.15, - center_y_fraction=-0.1, - edge_gain=0.3, - ) - case "Translation motion": - return TranslationMotionDistortion(shift_x_pixels=60, shift_y_pixels=10) - case "Rotational motion": - return RotationalMotionDistortion(angle_radians=torch.pi / 6) - case "Segmented rotational motion": - return SegmentedRotationalMotionDistortion( - angle_radians=(0.0, torch.pi / 20, -torch.pi / 24, torch.pi / 16), - ) - case "Segmented translation motion": - return SegmentedTranslationMotionDistortion( - shift_x_pixels=(0.0, 20.0, 50.0, -50.0), - shift_y_pixels=(0.0, 10.0, -20.0, 20.0), - ) - case "Gaussian bias field": - return GaussianKspaceBiasField(width_fraction=0.35, edge_gain=0.4) - case "Gaussian noise": - return GaussianNoiseDistortion(sigma=0.00001) - case "no distortion": - return BaseDistortion() - case _: - raise ValueError(f"Unknown distortion {name!r}") - - def choose_metric(name: str) -> dinv.metric.Metric: """Build one evaluation metric used in the saved comparison plots.""" @@ -258,7 +189,7 @@ def prepare_measurement_sample( x = sample_batch["x"].to(run_device) y = image_to_kspace(x) coil_maps = None - elif dataset_name in ("fastmri",) and use_oasis_fft_path: + elif dataset_name == "fastmri" and use_oasis_fft_path: y = sample_batch[1].to(run_device) x = fastmri_measurement_to_image(y) y = fastmri_measurement_to_oasis_kspace(y, device=run_device) @@ -431,7 +362,7 @@ def build_physics_pair( for algo_name in selected_algorithms: try: - use_oasis_path = uses_oasis_centered_path(args.dataset, algo_name) + use_oasis_path = uses_oasis_centered_path(algo_name) x_reference, y, coil_maps = prepare_measurement_sample( sample_batch=batch, dataset_name=args.dataset, @@ -446,75 +377,76 @@ def build_physics_pair( dataset=args.dataset, ).to(device) - for distortion_name in selected_distortions: - distortion = choose_distortion( - distortion_name, - keep_fraction=args.keep_fraction, - center_fraction=args.center_fraction, - cartesian_axis=-1 if use_oasis_path else -2, - ) - - physics_clean, physics = build_physics_pair( - image_shape=y.shape[-2:], - distortion_operator=distortion, - run_device=device, - use_oasis_fft_path=use_oasis_path, - coil_maps=coil_maps, - ) - y_distorted = distortion.A(y) - - # generate reference reconstructions (CG) for both clean and distorted k-space - # without correction for the distortion, i.e. using physics_clean in both cases - if use_oasis_path: - x_clean = x_reference - x_distorted = kspace_to_image(y_distorted) - else: - x_clean = ConjugateGradientReconstructor()(y, physics_clean) - x_distorted = ConjugateGradientReconstructor()(y_distorted, physics_clean) - - save_kspace_plot( - y, - y_distorted, - REPORT_DIR / f"DISTORTION_{algo_name}_{distortion_name}_sample_{i}.png", - distortion_name, - ) - - print( - f"Evaluating algo {algo_name}, distortion {distortion_name}, sample {i}..." - ) - - # actual reconstruction with the algo being evaluated - x_uncorrected = algo(y_distorted, physics_clean) - x_corrected = algo(y_distorted, physics) - - print("done!") - - dinv.utils.plot( - { - "Undistorted ksp, CG recon": x_clean, - "Distorted ksp, CG recon": x_distorted, - f"Distorted ksp, {algo_name} recon, uncorrected": x_uncorrected, - f"Distorted ksp, {algo_name} recon, corrected": x_corrected, - }, - subtitles=[ - "", - "", - "\n".join( - f"{m.__class__.__name__} {m(x_uncorrected, x_clean).item():.2f}" - for m in metrics - ), - "\n".join( - f"{m.__class__.__name__} {m(x_corrected, x_clean).item():.2f}" - for m in metrics - ), - ], - show=False, - close=True, - suptitle=f"Algo {algo_name}, distortion {distortion_name}, Sample {i}", - save_fn=REPORT_DIR / f"ALGO_{algo_name}_{distortion_name}_sample_{i}.png", - fontsize=3, - ) + for selected_distortion in selected_distortions: + for distortion_name, distortion_params in selected_distortion.items(): + distortion = choose_distortion_with_params( + distortion_name, + **distortion_params, + cartesian_axis=-1 if use_oasis_path else -2, + ) + + physics_clean, physics = build_physics_pair( + image_shape=y.shape[-2:], + distortion_operator=distortion, + run_device=device, + use_oasis_fft_path=use_oasis_path, + coil_maps=coil_maps, + ) + y_distorted = distortion.A(y) + + # generate reference reconstructions (CG) for both clean and distorted k-space + # without correction for the distortion, i.e. using physics_clean in both cases + if use_oasis_path: + x_clean = x_reference + x_distorted = kspace_to_image(y_distorted) + else: + x_clean = ConjugateGradientReconstructor()(y, physics_clean) + x_distorted = ConjugateGradientReconstructor()( + y_distorted, physics_clean + ) + + save_kspace_plot( + y, + y_distorted, + REPORT_DIR / f"DISTORTION_{algo_name}_{distortion_name}_sample_{i}.png", + distortion_name, + ) + + print( + f"Evaluating algo {algo_name}, distortion {distortion_name}, sample {i}..." + ) + + # actual reconstruction with the algo being evaluated + x_uncorrected = algo(y_distorted, physics_clean) + x_corrected = algo(y_distorted, physics) + + print("done!") + + dinv.utils.plot( + { + "Undistorted ksp, CG recon": x_clean, + "Distorted ksp, CG recon": x_distorted, + f"Distorted ksp, {algo_name} recon, uncorrected": x_uncorrected, + f"Distorted ksp, {algo_name} recon, corrected": x_corrected, + }, + subtitles=[ + "", + "", + "\n".join( + f"{m.__class__.__name__} {m(x_uncorrected, x_clean).item():.2f}" + for m in metrics + ), + "\n".join( + f"{m.__class__.__name__} {m(x_corrected, x_clean).item():.2f}" + for m in metrics + ), + ], + show=False, + close=True, + suptitle=f"Algo {algo_name}, distortion {distortion_name}, Sample {i}", + save_fn=REPORT_DIR + / f"ALGO_{algo_name}_{distortion_name}_sample_{i}.png", + fontsize=3, + ) except Exception as e: - print( - f"Error processing algo {algo_name}, distortion {distortion_name}, sample {i}: {e}" - ) + print(f"Error processing algo {algo_name}, sample {i}: {e}") diff --git a/examples/run_all.py b/examples/run_all.py index abebc8c..5230ca8 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -6,6 +6,10 @@ import os import sys +import glob +import SimpleITK as sitk +import numpy as np +import tqdm sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -14,12 +18,12 @@ import deepinv as dinv import torch import yaml -from tifffile import imwrite +from tifffile import imwrite, imread from mri_recon.distortions import ( BaseDistortion, DistortedKspaceMultiCoilMRI, - choose_distortion, + choose_distortion_with_params, ) from mri_recon.reconstruction import ( choose_reconstructor, @@ -100,10 +104,11 @@ def get_measurement_sample( # centered k-space data, shape: (B, 2, n_timepoints, (n_coils), H, W) dtype: float32 y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) - # select the first timepoint in order to simplify the evaluation of the reconstruction algorithms - y = y[:, :, 0, ...] - y_centered = y_centered[:, :, 0, ...] - x = x[:, :, 0, ...] + # select the center timepoint in order to simplify the evaluation of the reconstruction algorithms + center_time_point = y.shape[2] // 2 + y = y[:, :, center_time_point, ...] + y_centered = y_centered[:, :, center_time_point, ...] + x = x[:, :, center_time_point, ...] elif dataset_name == "fastmri_prostate": # reference image, shape: (B, W, H): dtype float32 @@ -120,15 +125,7 @@ def get_measurement_sample( return x, y, y_centered, coil_maps -if __name__ == "__main__": - # read config file in yaml format as first argument from commmand line - if len(sys.argv) < 2: - print("Usage: python examples/run_all.py ") - sys.exit(1) - - with open(sys.argv[1], "r") as f: - config = yaml.safe_load(f) - +def run_all(config) -> None: os.makedirs(config["results_dir"], exist_ok=True) # set up device @@ -192,177 +189,243 @@ def get_measurement_sample( BaseDistortion(), img_size=y.shape[-2:], coil_maps=coil_maps, device=device ) - # reference from dataset: - for distortion_name in config["distortions"]: - print(f"\t{distortion_name} ...") - - distortion = choose_distortion( - distortion_name, - keep_fraction=config["keep_fraction"], - center_fraction=config["center_fraction"], - cartesian_axis=-2, - ) - - y_distorted = distortion.A(y) - - physics_distorted = DistortedKspaceMultiCoilMRI( - distortion, - img_size=y.shape[-2:], - coil_maps=coil_maps, - device=device, - ) + for selected_distortion in config["distortions"]: + for distortion_name, distortion_params in selected_distortion.items(): + print(f"\t{distortion_name} ...") + distortion_name_with_params = distortion_name + "".join( + [p[0] + "=" + str(v) for p, v in distortion_params.items()] + ) + + distortion = choose_distortion_with_params( + distortion_name, + **distortion_params, + # keep_fraction=config["keep_fraction"], + # center_fraction=config["center_fraction"], + cartesian_axis=-2, + ) + + y_distorted = distortion.A(y) + + physics_distorted = DistortedKspaceMultiCoilMRI( + distortion, + img_size=y.shape[-2:], + coil_maps=coil_maps, + device=device, + ) + + for reconstructor_name in config["reconstruction_algorithms"]: + # only run on reconstructors, that use the fastmri-like k-space + if not uses_oasis_centered_path(reconstructor_name): + print(f"\t\t{reconstructor_name} ...") + start = datetime.now() + if compatible_dataset_with_reconstructor( + dataset_name, reconstructor_name + ): + reconstructor = choose_reconstructor( + reconstructor_name, + img_size=y_distorted.shape[-2:], + device=device, + verbose=config["verbose"], + ).to(device) + + # save reference and distorted k-space for debugging purposes + imwrite( + os.path.join( + config["results_dir"], + f"kspace_{dataset_name}_sample_{i}_reference.tiff", + ), + _kspace_to_log_magnitude(y).numpy(), + ) + imwrite( + os.path.join( + config["results_dir"], + f"kspace_{dataset_name}_sample_{i}_{distortion_name_with_params}.tiff", + ), + _kspace_to_log_magnitude(y_distorted).numpy(), + ) - for reconstructor_name in config["reconstruction_algorithms"]: - # only run on reconstructors, that use the fastmri-like k-space - if not uses_oasis_centered_path(reconstructor_name): - print(f"\t\t{reconstructor_name} ...") - start = datetime.now() - if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): - reconstructor = choose_reconstructor( - reconstructor_name, - img_size=y_distorted.shape[-2:], - device=device, - verbose=config["verbose"], - ).to(device) - - # save reference and distorted k-space for debugging purposes - imwrite( - os.path.join( - config["results_dir"], - f"kspace_{dataset_name}_sample_{i}_reference.tiff", - ), - _kspace_to_log_magnitude(y).numpy(), - ) - imwrite( - os.path.join( - config["results_dir"], - f"kspace_{dataset_name}_sample_{i}_{distortion_name}.tiff", - ), - _kspace_to_log_magnitude(y_distorted).numpy(), - ) + # actual reconstruction with the selected reconstructor + try: + x_uncorrected = reconstructor(y_distorted, physics_clean) + x_corrected = reconstructor(y_distorted, physics_distorted) + + # crop recostructed image to reference image size: + if x_uncorrected.shape[-2:] != x_reference.shape[-2:]: + x_uncorrected = physics_clean.crop( + x_uncorrected, shape=x_reference.shape[-2:] + ) + + if x_corrected.shape[-2:] != x_reference.shape[-2:]: + x_corrected = physics_distorted.crop( + x_corrected, shape=x_reference.shape[-2:] + ) + + # save reconstructed images + imwrite( + os.path.join( + config["results_dir"], + f"image_{dataset_name}_sample_{i}_{distortion_name_with_params}_{reconstructor_name}_uncorrected.tiff", + ), + convert_image_for_save(x_uncorrected), + ) + imwrite( + os.path.join( + config["results_dir"], + f"image_{dataset_name}_sample_{i}_{distortion_name_with_params}_{reconstructor_name}_corrected.tiff", + ), + convert_image_for_save(x_corrected), + ) + print(f"\t\t... done in {datetime.now() - start}") - # actual reconstruction with the selected reconstructor - try: - x_uncorrected = reconstructor(y_distorted, physics_clean) - x_corrected = reconstructor(y_distorted, physics_distorted) + except Exception as e: + print(f"Error using {reconstructor_name}: {e}") - # crop recostructed image to reference image size: - if x_uncorrected.shape[-2:] != x_reference.shape[-2:]: - x_uncorrected = physics_clean.crop( - x_uncorrected, shape=x_reference.shape[-2:] - ) + else: + print(f"\t\t ... not compatible with {dataset_name}") - if x_corrected.shape[-2:] != x_reference.shape[-2:]: - x_corrected = physics_distorted.crop( - x_corrected, shape=x_reference.shape[-2:] - ) + # now proceed with oasis-centered fft path + physics_clean = OasisCenteredFFTPhysics(BaseDistortion()) - # save reconstructed images + for selected_distortion in config["distortions"]: + for distortion_name, distortion_params in selected_distortion.items(): + print(f"\t{distortion_name} ...") + distortion_name_with_params = distortion_name + "".join( + [p[0] + "=" + str(v) for p, v in distortion_params.items()] + ) + + distortion = choose_distortion_with_params( + distortion_name, + **distortion_params, + cartesian_axis=-1, + ) + + y_distorted = torch.fft.fftshift( + distortion.A(torch.fft.fftshift(y_centered, dim=(-1, -2))), dim=(-2, -1) + ) + + physics_distorted = OasisCenteredFFTPhysics(distortion) + + for reconstructor_name in config["reconstruction_algorithms"]: + # skip all reconstructors, that don't use the oasis-centered path + if uses_oasis_centered_path(reconstructor_name): + print(f"\t\t{reconstructor_name} ...") + start = datetime.now() + if compatible_dataset_with_reconstructor( + dataset_name, reconstructor_name + ): + reconstructor = choose_reconstructor( + reconstructor_name, + img_size=y_distorted.shape[-2:], + device=device, + verbose=config["verbose"], + ).to(device) + + # save reference and distorted k-space for debugging purposes imwrite( os.path.join( config["results_dir"], - f"image_{dataset_name}_sample_{i}_{distortion_name}_{reconstructor_name}_uncorrected.tiff", + f"kspace_centered_{dataset_name}_sample_{i}_reference.tiff", ), - convert_image_for_save(x_uncorrected), + _kspace_to_log_magnitude(y_centered).numpy(), ) imwrite( os.path.join( config["results_dir"], - f"image_{dataset_name}_sample_{i}_{distortion_name}_{reconstructor_name}_corrected.tiff", + f"kspace_centered_{dataset_name}_sample_{i}_{distortion_name_with_params}.tiff", ), - convert_image_for_save(x_corrected), + _kspace_to_log_magnitude(y_distorted).numpy(), ) - print(f"\t\t... done in {datetime.now() - start}") - except Exception as e: - print(f"Error using {reconstructor_name}: {e}") + # actual reconstruction with the algo being evaluated + try: + x_uncorrected = reconstructor(y_distorted, physics_clean) + x_corrected = reconstructor(y_distorted, physics_distorted) + + if x_uncorrected.shape[-2:] != x_reference.shape[-2:]: + x_uncorrected = physics_clean.crop( + x_uncorrected, shape=x_reference.shape[-2:] + ) + + if x_corrected.shape[-2:] != x_reference.shape[-2:]: + x_corrected = physics_distorted.crop( + x_corrected, shape=x_reference.shape[-2:] + ) + + # save reconstructed images + imwrite( + os.path.join( + config["results_dir"], + f"image_{dataset_name}_sample_{i}_{distortion_name_with_params}_{reconstructor_name}_uncorrected.tiff", + ), + convert_image_for_save(x_uncorrected), + ) + imwrite( + os.path.join( + config["results_dir"], + f"image_{dataset_name}_sample_{i}_{distortion_name_with_params}_{reconstructor_name}_corrected.tiff", + ), + convert_image_for_save(x_corrected), + ) + print(f"\t\t... done in {datetime.now() - start}") - else: - print(f"\t\t ... not compatible with {dataset_name}") + except Exception as e: + print( + f"\t\tError using {reconstructor_name} with distortion {distortion_name_with_params} on sample {i}: {e}" + ) - # now proceed with oasis-centered fft path - physics_clean = OasisCenteredFFTPhysics(BaseDistortion()) + else: + print(f"\t\t ... not compatible with {dataset_name}") - for distortion_name in config["distortions"]: - print(f"\t{distortion_name} ...") - distortion = choose_distortion( - distortion_name, - keep_fraction=config["keep_fraction"], - center_fraction=config["center_fraction"], - cartesian_axis=-1, + if config["add_N4Correction"]: + reconstructed_images = glob.glob( + os.path.join(config["results_dir"], "*corrected.tiff") ) - - y_distorted = torch.fft.fftshift( - distortion.A(torch.fft.fftshift(y_centered, dim=(-1, -2))), dim=(-2, -1) + print( + f"Found {len(reconstructed_images)} in result folder, applying N4 Bias Field Correction" ) - - physics_distorted = OasisCenteredFFTPhysics(distortion) - - for reconstructor_name in config["reconstruction_algorithms"]: - # skip all reconstructors, that don't use the oasis-centered path - if uses_oasis_centered_path(reconstructor_name): - print(f"\t\t{reconstructor_name} ...") - start = datetime.now() - if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): - reconstructor = choose_reconstructor( - reconstructor_name, - img_size=y_distorted.shape[-2:], - device=device, - verbose=config["verbose"], - ).to(device) - - # save reference and distorted k-space for debugging purposes - imwrite( - os.path.join( - config["results_dir"], - f"kspace_centered_{dataset_name}_sample_{i}_reference.tiff", - ), - _kspace_to_log_magnitude(y_centered).numpy(), + with tqdm.tqdm( + total=len(reconstructed_images), desc="Applying N4 Bias Field Correction" + ) as pbar: + for reconstructed_image_filename in reconstructed_images: + reconstructed_image = imread(reconstructed_image_filename).squeeze() + if len(reconstructed_image.shape) == 2: + sitk_img = sitk.GetImageFromArray(reconstructed_image) + # sitk_mask = sitk.GetImageFromArray(mask.astype(np.uint8).T) + + corrector = sitk.N4BiasFieldCorrectionImageFilter() + corrector.SetMaximumNumberOfIterations([50] * 4) + corrector.SetConvergenceThreshold(0.001) + + # sitk.N4BiasFieldCorrectionImageFilter.Execute(corrector, sitk_img, sitk_mask) + sitk.N4BiasFieldCorrectionImageFilter.Execute(corrector, sitk_img) + + log_bias_sitk = corrector.GetLogBiasFieldAsImage(sitk_img) + bias_n4 = np.exp(sitk.GetArrayFromImage(log_bias_sitk).T) + reconstructed_image_n4 = reconstructed_image / np.where( + bias_n4 > 0, bias_n4, 1.0 ) imwrite( - os.path.join( - config["results_dir"], - f"kspace_centered_{dataset_name}_sample_{i}_{distortion_name}.tiff", + reconstructed_image_filename.replace( + "corrected.tiff", "corrected_N4.tiff" ), - _kspace_to_log_magnitude(y_distorted).numpy(), + reconstructed_image_n4, ) - # actual reconstruction with the algo being evaluated - try: - x_uncorrected = reconstructor(y_distorted, physics_clean) - x_corrected = reconstructor(y_distorted, physics_distorted) + else: + print("Skipping N4 Bias Field correction for image") + print(os.path.basename(reconstructed_image_filename)) + print("which as shape ", reconstructed_image.shape) - if x_uncorrected.shape[-2:] != x_reference.shape[-2:]: - x_uncorrected = physics_clean.crop( - x_uncorrected, shape=x_reference.shape[-2:] - ) + pbar.update(1) - if x_corrected.shape[-2:] != x_reference.shape[-2:]: - x_corrected = physics_distorted.crop( - x_corrected, shape=x_reference.shape[-2:] - ) - # save reconstructed images - imwrite( - os.path.join( - config["results_dir"], - f"image_{dataset_name}_sample_{i}_{distortion_name}_{reconstructor_name}_uncorrected.tiff", - ), - convert_image_for_save(x_uncorrected), - ) - imwrite( - os.path.join( - config["results_dir"], - f"image_{dataset_name}_sample_{i}_{distortion_name}_{reconstructor_name}_corrected.tiff", - ), - convert_image_for_save(x_corrected), - ) - print(f"\t\t... done in {datetime.now() - start}") +if __name__ == "__main__": + # read config file in yaml format as first argument from commmand line + if len(sys.argv) < 2: + print("Usage: python examples/run_all.py ") + sys.exit(1) - except Exception as e: - print( - f"\t\tError using {reconstructor_name} with distortion {distortion_name} on sample {i}: {e}" - ) + with open(sys.argv[1], "r") as f: + config = yaml.safe_load(f) - else: - print(f"\t\t ... not compatible with {dataset_name}") + run_all(config) diff --git a/mri_recon/distortions/__init__.py b/mri_recon/distortions/__init__.py index ea85943..874b248 100644 --- a/mri_recon/distortions/__init__.py +++ b/mri_recon/distortions/__init__.py @@ -3,7 +3,11 @@ DistortedKspaceMultiCoilMRI, SelfAdjointMultiplicativeMaskDistortion, ) -from .biasfield import GaussianKspaceBiasField, OffCenterAnisotropicGaussianKspaceBiasField +from .biasfield import ( + GaussianKspaceBiasField, + GaussianBiasField, + OffCenterAnisotropicGaussianKspaceBiasField, +) from .ghosting import PhaseEncodeGhostingDistortion from .motion import ( RotationalMotionDistortion, @@ -20,4 +24,4 @@ RadialHighPassEmphasisDistortion, ) from .undersampling import CartesianUndersampling, PartialFourierDistortion -from .utils import choose_distortion +from .utils import choose_distortion_with_params diff --git a/mri_recon/distortions/base.py b/mri_recon/distortions/base.py index c5b0734..80dae2f 100644 --- a/mri_recon/distortions/base.py +++ b/mri_recon/distortions/base.py @@ -75,6 +75,50 @@ def _validate_cartesian_kspace_tensor(y: torch.Tensor) -> None: raise ValueError(f"Spatial k-space dimensions must be positive, got shape {tuple(y.shape)}") +def shifted_kspace_to_image(y: torch.Tensor) -> torch.Tensor: + """Convert centered channel-first k-space to complex images. + + Parameters + ---------- + y : torch.Tensor + Centered k-space tensor with shape ``(B, 2, H, W)``. + + Returns + ------- + torch.Tensor + Complex image tensor with shape ``(B, 2, H, W)``. + """ + + y_complex = torch.view_as_complex(y.movedim(1, -1).contiguous()) + x_complex = torch.fft.fftshift( + torch.fft.ifft2(torch.fft.ifftshift(y_complex, dim=(-2, -1)), dim=(-2, -1), norm="ortho"), + dim=(-2, -1), + ) + return torch.view_as_real(x_complex).movedim(-1, 1).contiguous() + + +def image_to_shifted_kspace(x: torch.Tensor) -> torch.Tensor: + """Convert channel-first complex images to centered k-space. + + Parameters + ---------- + x : torch.Tensor + Complex image tensor with shape ``(B, 2, H, W)``. + + Returns + ------- + torch.Tensor + Centered k-space tensor with shape ``(B, 2, H, W)``. + """ + + x_complex = torch.view_as_complex(x.movedim(1, -1).contiguous()) + y_complex = torch.fft.fftshift( + torch.fft.fft2(torch.fft.ifftshift(x_complex, dim=(-2, -1)), dim=(-2, -1), norm="ortho"), + dim=(-2, -1), + ) + return torch.view_as_real(y_complex).movedim(-1, 1).contiguous() + + class BaseDistortion(dinv.physics.LinearPhysics): """Base class for deterministic k-space distortions. diff --git a/mri_recon/distortions/biasfield.py b/mri_recon/distortions/biasfield.py index a9d7910..9136c9c 100644 --- a/mri_recon/distortions/biasfield.py +++ b/mri_recon/distortions/biasfield.py @@ -4,7 +4,12 @@ import torch -from mri_recon.distortions.base import BaseDistortion, _normalized_frequency_grids +from mri_recon.distortions.base import ( + BaseDistortion, + _normalized_frequency_grids, + image_to_shifted_kspace, + shifted_kspace_to_image, +) def _gaussian_bias_gain_field( @@ -34,6 +39,36 @@ def _gaussian_bias_gain_field( return gain / gain.max() +def _gaussian_bias_field( + shape: tuple[int, ...], + device: torch.device, + width_x_fraction: float, + width_y_fraction: float, + center_x_fraction: float, + center_y_fraction: float, + edge_gain: float, +) -> torch.Tensor: + """Build a normalized Gaussian multiplicative bias field in image space.""" + + height = shape[-2] + width = shape[-1] + ky = torch.linspace(-1.0, 1.0, steps=height, device=device) + kx = torch.linspace(-1.0, 1.0, steps=width, device=device) + ky_grid, kx_grid = torch.meshgrid(ky, kx, indexing="ij") + + dx = (kx_grid - center_x_fraction) / width_x_fraction + dy = (ky_grid - center_y_fraction) / width_y_fraction + gaussian = torch.exp(-0.5 * (dx * dx + dy * dy)) + + edge_value = float(gaussian.min()) + if edge_value < 1.0: + gaussian = (gaussian - edge_value) / (1.0 - edge_value) + gaussian = gaussian.clamp(0.0, 1.0) + + gain = edge_gain + (1.0 - edge_gain) * gaussian + return gain / gain.max() + + class GaussianKspaceBiasField(BaseDistortion): """Smooth centered multiplicative bias field in k-space. @@ -74,6 +109,47 @@ def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: return self.A(y) +class GaussianBiasField(BaseDistortion): + """Smooth centered multiplicative bias field in image space. + + The gain is radial, equals ``1`` at DC, and smoothly decays toward + ``edge_gain`` at the edge of the sampled k-space grid. + + :param float width_fraction: Radial width of the Gaussian envelope on the + normalized k-space grid. + :param float edge_gain: Gain approached near the edge of image space. Must lie + in ``(0, 1]``. + """ + + def __init__(self, width_fraction: float = 0.35, edge_gain: float = 0.4) -> None: + super().__init__() + if width_fraction <= 0.0: + raise ValueError("width_fraction must be positive") + if not 0.0 < edge_gain <= 1.0: + raise ValueError("edge_gain must be in (0, 1]") + self.width_fraction = width_fraction + self.edge_gain = edge_gain + + def _gain_field(self, shape: tuple[int, ...], device: torch.device) -> torch.Tensor: + return _gaussian_bias_field( + shape=shape, + device=device, + width_x_fraction=self.width_fraction, + width_y_fraction=self.width_fraction, + center_x_fraction=0.0, + center_y_fraction=0.0, + edge_gain=self.edge_gain, + ) + + def A(self, y: torch.Tensor) -> torch.Tensor: + gain = self._gain_field(y.shape, y.device) + B = 1.0 / (torch.exp(gain)) + return image_to_shifted_kspace(shifted_kspace_to_image(y) * B) + + def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: + return self.A(y) + + class OffCenterAnisotropicGaussianKspaceBiasField(BaseDistortion): """Off-center anisotropic Gaussian multiplicative bias field in k-space. @@ -151,3 +227,57 @@ def A(self, y: torch.Tensor) -> torch.Tensor: def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: return self.A(y) + + +class OffCenterAnisotropicGaussianBiasField(BaseDistortion): + """Off-center anisotropic Gaussian multiplicative bias field . + + + + :param float width_x_fraction: + :param float width_y_fraction: + :param float center_x_fraction: + :param float center_y_fraction: + :param float edge_gain: Baseline gain far from the Gaussian peak. Must lie + in ``(0, 1]``. Smaller values strengthen the peripheral attenuation. + """ + + def __init__( + self, + width_x_fraction: float = 0.2, + width_y_fraction: float = 0.35, + center_x_fraction: float = 0.15, + center_y_fraction: float = -0.1, + edge_gain: float = 0.3, + ) -> None: + super().__init__() + if width_x_fraction <= 0.0 or width_y_fraction <= 0.0: + raise ValueError("width fractions must be positive") + if abs(center_x_fraction) > 1.0 or abs(center_y_fraction) > 1.0: + raise ValueError("center fractions must be in [-1, 1]") + if not 0.0 < edge_gain <= 1.0: + raise ValueError("edge_gain must be in (0, 1]") + self.width_x_fraction = width_x_fraction + self.width_y_fraction = width_y_fraction + self.center_x_fraction = center_x_fraction + self.center_y_fraction = center_y_fraction + self.edge_gain = edge_gain + + def _gain_field(self, shape: tuple[int, ...], device: torch.device) -> torch.Tensor: + return _gaussian_bias_field( + shape=shape, + device=device, + width_x_fraction=self.width_x_fraction, + width_y_fraction=self.width_y_fraction, + center_x_fraction=self.center_x_fraction, + center_y_fraction=self.center_y_fraction, + edge_gain=self.edge_gain, + ) + + def A(self, y: torch.Tensor) -> torch.Tensor: + gain = self._gain_field(y.shape, y.device) + B = 1.0 / (torch.exp(gain)) + return image_to_shifted_kspace(shifted_kspace_to_image(y) * B) + + def A_adjoint(self, y: torch.Tensor) -> torch.Tensor: + return self.A(y) diff --git a/mri_recon/distortions/ghosting.py b/mri_recon/distortions/ghosting.py index 9ba07c0..a62a97d 100644 --- a/mri_recon/distortions/ghosting.py +++ b/mri_recon/distortions/ghosting.py @@ -36,7 +36,8 @@ def __init__( self, line_period: int = 2, line_offset: int = 1, - phase_error_radians: float = torch.pi / 2, + phase_error_radians: float | None = None, + phase_error_degrees: float | None = None, corrupted_line_scale: float = 1.0, ghost_axis: int = -2, ) -> None: @@ -52,6 +53,11 @@ def __init__( self.line_period = int(line_period) self.line_offset = int(line_offset) + if phase_error_radians is None: + if phase_error_degrees is None: + raise ValueError("phase_error_radians or phase_error_degrees must not be None") + else: + phase_error_radians = 2 * torch.pi * phase_error_degrees / 360.0 self.phase_error_radians = float(phase_error_radians) self.corrupted_line_scale = float(corrupted_line_scale) self.ghost_axis = ghost_axis diff --git a/mri_recon/distortions/motion.py b/mri_recon/distortions/motion.py index e4f4fe6..d8d0f0c 100644 --- a/mri_recon/distortions/motion.py +++ b/mri_recon/distortions/motion.py @@ -65,8 +65,16 @@ class RotationalMotionDistortion(BaseDistortion): :param float angle_radians: In-plane rotation angle in radians. """ - def __init__(self, angle_radians: float = torch.pi / 12) -> None: + def __init__( + self, angle_radians: float | None = None, angle_degrees: float | None = None + ) -> None: super().__init__() + if angle_radians is None: + if angle_degrees is None: + angle_radians = torch.pi / 12 + else: + angle_radians = 2 * torch.pi * angle_degrees / 360.0 + self.angle_radians = float(angle_radians) def _reshape_kspace_channels(self, y: torch.Tensor) -> tuple[torch.Tensor, tuple[int, ...]]: @@ -192,11 +200,27 @@ class SegmentedRotationalMotionDistortion(BaseDistortion): motion changes across the phase-encode lines. """ - def __init__(self, angle_radians: tuple[float, ...], segment_axis: int = -2) -> None: + def __init__( + self, + angle_radians: tuple[float, ...] | None = None, + angle_degrees: tuple[float, ...] | None = None, + segment_axis: int = -2, + ) -> None: super().__init__() - if len(angle_radians) == 0: + + if angle_radians is None: + if angle_degrees is None: + raise ValueError("Either angle_radians or angle_degrees must not be None") + else: + if len(angle_degrees) == 0: + raise ValueError("angle_degrees must be non-empty list") + else: + angle_radians = [2 * torch.pi * alpha / 360.0 for alpha in angle_degrees] + elif len(angle_radians) == 0: raise ValueError("angle_radians must be non-empty") + self.angle_radians = angle_radians + if segment_axis not in (-2, -1): raise ValueError("segment_axis must be -2 or -1 for 2D k-space") diff --git a/mri_recon/distortions/utils.py b/mri_recon/distortions/utils.py index 96d00b7..fe6157e 100644 --- a/mri_recon/distortions/utils.py +++ b/mri_recon/distortions/utils.py @@ -9,7 +9,12 @@ RadialHighPassEmphasisDistortion, ) from .undersampling import CartesianUndersampling, PartialFourierDistortion -from .biasfield import OffCenterAnisotropicGaussianKspaceBiasField, GaussianKspaceBiasField +from .biasfield import ( + OffCenterAnisotropicGaussianKspaceBiasField, + GaussianKspaceBiasField, + GaussianBiasField, + OffCenterAnisotropicGaussianBiasField, +) from .noise import GaussianNoiseDistortion from .motion import ( RotationalMotionDistortion, @@ -116,6 +121,14 @@ def choose_distortion( center_y_fraction=-0.1, edge_gain=0.3, ) + case "OffCenterAnisotropicGaussianBiasField": + return OffCenterAnisotropicGaussianBiasField( + width_x_fraction=0.2, + width_y_fraction=0.35, + center_x_fraction=0.15, + center_y_fraction=-0.1, + edge_gain=0.3, + ) case "TranslationMotion": return TranslationMotionDistortion(shift_x_pixels=60, shift_y_pixels=10) case "RotationalMotion": @@ -131,6 +144,8 @@ def choose_distortion( ) case "GaussianKspaceBiasField": return GaussianKspaceBiasField(width_fraction=0.35, edge_gain=0.4) + case "GaussianBiasField": + return GaussianKspaceBiasField(width_fraction=0.35, edge_gain=0.4) case "GaussianNoise": return GaussianNoiseDistortion(sigma=0.00001) case "BaseDistortion": @@ -138,3 +153,75 @@ def choose_distortion( case _: raise ValueError(f"Unknown distortion {name!r}") + + +def choose_distortion_with_params(name: str, cartesian_axis: int = -2, **kwargs) -> BaseDistortion: + """Build one distortion operator for the inference comparison script. + + The ``cartesian_axis`` is supplied by the active measurement convention: + FastMRI-native runs use the repository's existing axis, while OASIS-native + and FastMRI-to-OASIS runs use the centered OASIS axis. + """ + + match name: + case "PhaseEncodeGhosting": + return PhaseEncodeGhostingDistortion(**kwargs) + case "CartesianUndersamplingVariableDensity": + return CartesianUndersampling( + **kwargs, axis=cartesian_axis, pattern="variable_density_random", seed=42 + ) + case "CartesianUndersamplingUniformRandom": + return CartesianUndersampling( + **kwargs, axis=cartesian_axis, pattern="uniform_random", seed=42 + ) + case "CartesianUndersamplingUniformRandomZeroACS": + return CartesianUndersampling( + **kwargs, + axis=cartesian_axis, + pattern="uniform_random", + seed=42, + center_fraction=0.0, + ) + case "CartesianUndersamplingEquispaced": + return CartesianUndersampling( + **kwargs, axis=cartesian_axis, pattern="equispaced", seed=42 + ) + case "CartesianUndersamplingEquispacedZeroACS": + return CartesianUndersampling( + **kwargs, axis=cartesian_axis, pattern="equispaced", seed=42, center_fraction=0.0 + ) + case "PartialFourier": + return PartialFourierDistortion(**kwargs, axis=cartesian_axis) + case "AnisotropicLP": + return AnisotropicResolutionReduction(**kwargs) + case "HannTaperLP": + return HannTaperResolutionReduction(**kwargs) + case "KaiserTaperLP": + return KaiserTaperResolutionReduction(**kwargs) + case "RadialHighPassEmphasis": + return RadialHighPassEmphasisDistortion(**kwargs) + case "IsotropicLP": + return IsotropicResolutionReduction(**kwargs) + case "OffCenterAnisotropicGaussianKspaceBiasField": + return OffCenterAnisotropicGaussianKspaceBiasField(**kwargs) + case "OffCenterAnisotropicGaussianBiasField": + return OffCenterAnisotropicGaussianBiasField(**kwargs) + case "TranslationMotion": + return TranslationMotionDistortion(**kwargs) + case "RotationalMotion": + return RotationalMotionDistortion(**kwargs) + case "SegmentedRotationalMotion": + return SegmentedRotationalMotionDistortion(**kwargs) + case "SegmentedTranslationMotion": + return SegmentedTranslationMotionDistortion(**kwargs) + case "GaussianKspaceBiasField": + return GaussianKspaceBiasField(**kwargs) + case "GaussianBiasField": + return GaussianBiasField(**kwargs) + case "GaussianNoise": + return GaussianNoiseDistortion(**kwargs) + case "BaseDistortion": + return BaseDistortion() + + case _: + raise ValueError(f"Unknown distortion {name!r}") diff --git a/pyproject.toml b/pyproject.toml index c619e37..4a45111 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "pytest>=9.0.2", "python-certifi-win32>=1.6.1", "sigpy>=0.1.27", + "SimpleITK>=2.5.5", "torch>=2.11.0", "torchmetrics>=1.9.0", "torchvision>=0.26.0", diff --git a/uv.lock b/uv.lock index 30e3765..1b5421b 100644 --- a/uv.lock +++ b/uv.lock @@ -139,6 +139,7 @@ dependencies = [ { name = "pytest" }, { name = "python-certifi-win32" }, { name = "sigpy" }, + { name = "simpleitk" }, { name = "torch", version = "2.11.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" }, { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -170,6 +171,7 @@ requires-dist = [ { name = "pytest", specifier = ">=9.0.2" }, { name = "python-certifi-win32", specifier = ">=1.6.1" }, { name = "sigpy", specifier = ">=0.1.27" }, + { name = "simpleitk", specifier = ">=2.5.5" }, { name = "torch", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", specifier = ">=2.11.0" }, { name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.11.0", index = "https://download.pytorch.org/whl/cpu" }, { name = "torch", marker = "sys_platform == 'linux'", specifier = ">=2.11.0", index = "https://download.pytorch.org/whl/cu128" }, @@ -2063,6 +2065,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/5d/e766856bba1593dafbbf718112969c3b656d79d4e943ca5dbfa2e4453e9c/sigpy-0.1.27-py3-none-any.whl", hash = "sha256:c0b2f2039ff2dc4497890f615e417860a14e40ee2c22b1fb2571ad7383e1c2c6", size = 1957, upload-time = "2025-01-10T21:28:13.38Z" }, ] +[[package]] +name = "simpleitk" +version = "2.5.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/28/8f/c21ee6dc4dbede4fc385b5d70703d45f749db2738c7b0c175546f43c9f86/simpleitk-2.5.5.tar.gz", hash = "sha256:254f70febed55868801c89f310e650eb1b87148d82b1d4e3628d8f0181a7f529", size = 2112106, upload-time = "2026-05-13T20:28:15.503Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/79/39/b494f7c90a2739369ef04f7f09c67d29a8dd3150fff51b0c2f89111b7f81/simpleitk-2.5.5-cp311-abi3-macosx_10_9_x86_64.whl", hash = "sha256:e831d15f5ea5b0e02ed3a6ac5f2dd34ba0a0eeaa4ebacfb1d970999836b19ee0", size = 42679794, upload-time = "2026-05-13T20:24:49.959Z" }, + { url = "https://files.pythonhosted.org/packages/6e/34/4b3208b35dea488263a5c9f4a464ef20316f663e9e90d5de61349c31b327/simpleitk-2.5.5-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:ee87416622cec6ed6b96747a0648c82c6e88f10b6553668ce90736f09ab5a994", size = 38646487, upload-time = "2026-05-13T20:25:09.636Z" }, + { url = "https://files.pythonhosted.org/packages/14/e5/6a77865140278c86a392a85bd2317a1097bd891f54955ebe5db623e4b37c/simpleitk-2.5.5-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:847f42057d7bf01b5721b107dd3e00d4ee8d514f6f1ff85794ffedacaf4afc53", size = 48082097, upload-time = "2026-05-13T20:25:26.946Z" }, + { url = "https://files.pythonhosted.org/packages/9f/68/ed67a355a62848ee04bb4f01e89d3be871052c2c3ae6d5fc0fb2f6010979/simpleitk-2.5.5-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8ace5e392be00d5bec28c87ef1d73016f8ebe671dafffe9e9d045f20b0968033", size = 52778654, upload-time = "2026-05-13T20:25:41.737Z" }, + { url = "https://files.pythonhosted.org/packages/42/ab/3a03de43173749171340ad974636105cb02c780f76055551c874c3df7fff/simpleitk-2.5.5-cp311-abi3-win_amd64.whl", hash = "sha256:a5fcfcfe9242d3d509b254b3213ca0f5db2c15903f2fc375d9ac5d38e57d415d", size = 18920455, upload-time = "2026-05-13T20:25:48.765Z" }, + { url = "https://files.pythonhosted.org/packages/58/e9/e2ea5fc95dbc66dbed925634ed308ea5246d1ec923dfdf74cc82e081acd4/simpleitk-2.5.5-cp314-cp314t-macosx_10_9_x86_64.whl", hash = "sha256:2e486cac9d0a147c70ebf983e7e5fb49da1d039f1a3393a9b04dcdfdc78899b4", size = 42709520, upload-time = "2026-05-13T20:26:09.907Z" }, + { url = "https://files.pythonhosted.org/packages/82/ac/8208b28cf1b3a998ab4d730f8f80acfcffd2fc023a5f7e4090c4b4ff4c38/simpleitk-2.5.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a45327a014c730c4bf903de288162a608857df7b9f91ace39e8600c6f8d00465", size = 38662876, upload-time = "2026-05-13T20:26:25.906Z" }, + { url = "https://files.pythonhosted.org/packages/74/15/83d370f2b37ab224af0a9a5e83511f55046a1f1f51395c552916443cefe7/simpleitk-2.5.5-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bd0d4552648dfce62f7462116a131ed9240fe072fb9abaa1efa31464b637937c", size = 47935456, upload-time = "2026-05-13T20:26:57.97Z" }, + { url = "https://files.pythonhosted.org/packages/9d/bd/4f4998f4c5f581282d9cedd0e965e0e8b6fb0f1e3ca9113c9cffcbe22b72/simpleitk-2.5.5-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:29c4e26fe02c4c0d67bbdd89225ad655eed0b8d6346b875aa9b09018773009c2", size = 52629787, upload-time = "2026-05-13T20:27:14.4Z" }, + { url = "https://files.pythonhosted.org/packages/dd/eb/0df2167bf35f502c90576ef39db58e0728159429bbb0fdccf35acaae63b7/simpleitk-2.5.5-cp314-cp314t-win_amd64.whl", hash = "sha256:dd204e3490014cd475a654817ca9d593e69c9b5f637290478fd0020c0f84406b", size = 19514461, upload-time = "2026-05-13T20:27:21.388Z" }, +] + [[package]] name = "six" version = "1.17.0" From c17c382c5d549fe80e5371022f985938bb903ea9 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Tue, 9 Jun 2026 15:23:10 +0000 Subject: [PATCH 13/22] crop reconstructed images to reference image size, correct dimensions, when image to shifted k-space --- examples/config.yaml | 24 +++++++--- examples/fastmri_inference_plot.py | 74 ++++++++++-------------------- examples/run_all.py | 34 ++++++++------ mri_recon/distortions/__init__.py | 2 + 4 files changed, 65 insertions(+), 69 deletions(-) diff --git a/examples/config.yaml b/examples/config.yaml index 65f6d84..304002b 100644 --- a/examples/config.yaml +++ b/examples/config.yaml @@ -39,14 +39,23 @@ distortions: - "RotationalMotion": # angle_radians=torch.pi / 6 "angle_degrees": 60.0 + # - "OffCenterAnisotropicGaussianKspaceBiasField": + # "width_x_fraction": 0.2 + # "width_y_fraction": 0.35 + # "center_x_fraction": 0.15 + # "center_y_fraction": -0.1 + # "edge_gain": 0.3 - "OffCenterAnisotropicGaussianBiasField": "width_x_fraction": 0.2 "width_y_fraction": 0.35 "center_x_fraction": 0.15 "center_y_fraction": -0.1 "edge_gain": 0.05 + # - "GaussianKspaceBiasField": + # "width_fraction": 0.35 + # "edge_gain": 0.4 - "GaussianBiasField": - "width_fraction": 0.35 + "width_fraction": 0.15 "edge_gain": 0.05 - "AnisotropicLP": "kx_radius_fraction": 1.0, @@ -67,16 +76,17 @@ distortions: reconstruction_algorithms: - "zero-filled" - "conjugate-gradient" - #- "ram", - #- "dip", - #- "tv-pgd", - #- "wavelet-fista", - #- "tv-fista", - #- "tv-pdhg", + #- "ram" + #- "dip" + #- "tv-pgd" + #- "wavelet-fista" + #- "tv-fista" + #- "tv-pdhg" - "unet-fastmri" - "unet-oasis-acceleration4" #- "unet-oasis-acceleration8" #- "unet-oasis-acceleration10" +# additionally performs N4 Bias Field correction on reference images and BiasField-distorted and reconstructed ones add_N4Correction: false num_samples: 1 verbose: true diff --git a/examples/fastmri_inference_plot.py b/examples/fastmri_inference_plot.py index 9b34fd8..d9f73ac 100644 --- a/examples/fastmri_inference_plot.py +++ b/examples/fastmri_inference_plot.py @@ -18,6 +18,7 @@ choose_distortion_with_params, BaseDistortion, DistortedKspaceMultiCoilMRI, + image_to_shifted_kspace, ) from mri_recon.reconstruction import ( ConjugateGradientReconstructor, @@ -30,21 +31,11 @@ from mri_recon.utils import ( OasisCenteredFFTPhysics, OasisSliceDataset, - fastmri_measurement_to_image, fastmri_measurement_to_oasis_kspace, - image_to_kspace, kspace_to_image, save_kspace_plot, ) -FASTMRI_REPORT_DIR = Path("reports") / "fastmri_inference_plot" -FASTMRI_MULTICOIL_REPORT_DIR = Path("reports") / "fastmri_multicoil_inference_plot" -OASIS_REPORT_DIR = Path("reports") / "oasis_inference_plot" -CMRXRECON_REPORT_DIR = Path("reports") / "cmrxrecon_inference_plot" -FASTMRI_REPORT_DIR.mkdir(parents=True, exist_ok=True) -FASTMRI_MULTICOIL_REPORT_DIR.mkdir(parents=True, exist_ok=True) -OASIS_REPORT_DIR.mkdir(parents=True, exist_ok=True) -CMRXRECON_REPORT_DIR.mkdir(parents=True, exist_ok=True) ALGORITHMS = [ # "zero-filled", "conjugate-gradient", @@ -187,27 +178,11 @@ def prepare_measurement_sample( if dataset_name == "oasis": x = sample_batch["x"].to(run_device) - y = image_to_kspace(x) + y = image_to_shifted_kspace(x) coil_maps = None - elif dataset_name == "fastmri" and use_oasis_fft_path: - y = sample_batch[1].to(run_device) - x = fastmri_measurement_to_image(y) - y = fastmri_measurement_to_oasis_kspace(y, device=run_device) - coil_maps = None - elif dataset_name == "fastmri_multicoil" and use_oasis_fft_path: - y = sample_batch[1].to(run_device) - coil_maps = ( - sample_batch[2]["coil_maps"].to(run_device) - if isinstance(sample_batch, (tuple, list)) - and len(sample_batch) == 3 - and "coil_maps" in sample_batch[2] - else None - ) - x = fastmri_measurement_to_image(y, coil_maps=coil_maps) - y = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) elif dataset_name in ("fastmri", "fastmri_multicoil"): - x = None + x = sample_batch[0].to(run_device) y = sample_batch[1].to(run_device) coil_maps = ( sample_batch[2]["coil_maps"].to(run_device) @@ -227,6 +202,9 @@ def prepare_measurement_sample( else None ) + if use_oasis_fft_path: + y = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) + return x, y, coil_maps @@ -275,18 +253,6 @@ def build_physics_pair( ) parser.add_argument("--distortion", type=str, default="", choices=DISTORTIONS) - parser.add_argument( - "--keep_fraction", - type=float, - default=0.25, - help="Fraction of k-space lines to keep for undersampling distortions.", - ) - parser.add_argument( - "--center_fraction", - type=float, - default=0.125, - help="Fraction of low-frequency k-space lines to keep fully for undersampling distortions.", - ) # algo related arguments parser.add_argument( @@ -308,6 +274,9 @@ def build_physics_pair( selected_algorithms = ALGORITHMS if args.algorithm == "" else [args.algorithm] selected_distortions = DISTORTIONS if args.distortion == "" else [args.distortion] + REPORT_DIR = Path("reports") / Path(args.dataset + "_inference_plot") + REPORT_DIR.mkdir(parents=True, exist_ok=True) + # skip non-compatible algorithm-dataset pairs selected_algorithms = [ algo_name @@ -316,15 +285,7 @@ def build_physics_pair( ] # set up report dir - if args.dataset == "fastmri": - REPORT_DIR = FASTMRI_REPORT_DIR - elif args.dataset == "oasis": - REPORT_DIR = OASIS_REPORT_DIR - elif args.dataset == "fastmri_multicoil": - REPORT_DIR = FASTMRI_MULTICOIL_REPORT_DIR - elif args.dataset == "cmrxrecon": - REPORT_DIR = CMRXRECON_REPORT_DIR - else: + if args.dataset not in ["fastmri", "oasis", "fastmri_multicoil", "cmrxrecon"]: raise NotImplementedError(f"Invalid dataset: {args.dataset}") # set up device, dataset, metrics @@ -405,6 +366,12 @@ def build_physics_pair( y_distorted, physics_clean ) + if x_clean.shape[-2:] != x_reference.shape[-2:]: + x_clean = physics_clean.crop(x_clean, shape=x_reference.shape[-2:]) + + if x_distorted.shape[-2:] != x_reference.shape[-2:]: + x_distorted = physics.crop(x_distorted, shape=x_reference.shape[-2:]) + save_kspace_plot( y, y_distorted, @@ -418,7 +385,16 @@ def build_physics_pair( # actual reconstruction with the algo being evaluated x_uncorrected = algo(y_distorted, physics_clean) + + # crop recostructed image to reference image size: + if x_uncorrected.shape[-2:] != x_reference.shape[-2:]: + x_uncorrected = physics_clean.crop( + x_uncorrected, shape=x_reference.shape[-2:] + ) + x_corrected = algo(y_distorted, physics) + if x_corrected.shape[-2:] != x_reference.shape[-2:]: + x_corrected = physics.crop(x_corrected, shape=x_reference.shape[-2:]) print("done!") diff --git a/examples/run_all.py b/examples/run_all.py index 5230ca8..e58833f 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -24,6 +24,7 @@ BaseDistortion, DistortedKspaceMultiCoilMRI, choose_distortion_with_params, + image_to_shifted_kspace, ) from mri_recon.reconstruction import ( choose_reconstructor, @@ -35,7 +36,6 @@ OasisCenterSliceFolderDataset, FastMRIProstateDataset, fastmri_measurement_to_oasis_kspace, - oasis_kspace_to_fastmri_measurement, image_to_kspace, _kspace_to_log_magnitude, convert_image_for_save, @@ -60,7 +60,8 @@ def get_measurement_sample( # centered k-space data, shape: (B, 2, H, W) dtype: float32 y_centered = image_to_kspace(x) # k-space data, shape: (B, 2, H, W) dtype: float32 - y = oasis_kspace_to_fastmri_measurement(y_centered) + # y = oasis_kspace_to_fastmri_measurement(y_centered) + y = image_to_shifted_kspace(x) elif dataset_name == "fastmri_knee": # reference image, shape: (B, 1, H/2, H/2) dtype: float32 x = sample_batch[0].to(run_device) @@ -120,7 +121,14 @@ def get_measurement_sample( # (B, 2, H, W) y_centered = image_to_kspace(x) - y = oasis_kspace_to_fastmri_measurement(y_centered) + # y = oasis_kspace_to_fastmri_measurement(y_centered) + y = image_to_shifted_kspace(x) + + print("Debug shapes:") + print("x: ", x.shape) + print("y: ", y.shape) + print("y_centered: ", y_centered.shape) + print("coil_maps: ", coil_maps.shape if coil_maps is not None else "None") return x, y, y_centered, coil_maps @@ -377,19 +385,21 @@ def run_all(config) -> None: print(f"\t\t ... not compatible with {dataset_name}") if config["add_N4Correction"]: - reconstructed_images = glob.glob( - os.path.join(config["results_dir"], "*corrected.tiff") + reconstructed_bias_field_images = glob.glob( + os.path.join(config["results_dir"], "*BiasField*corrected.tiff") ) - print( - f"Found {len(reconstructed_images)} in result folder, applying N4 Bias Field Correction" + reference_images = glob.glob( + os.path.join(config["results_dir"], "image*reference.tiff") ) + images_for_n4_correction = reconstructed_bias_field_images + reference_images with tqdm.tqdm( - total=len(reconstructed_images), desc="Applying N4 Bias Field Correction" + total=len(images_for_n4_correction), desc="Applying N4 Bias Field Correction" ) as pbar: - for reconstructed_image_filename in reconstructed_images: + for reconstructed_image_filename in images_for_n4_correction: + print(reconstructed_image_filename) reconstructed_image = imread(reconstructed_image_filename).squeeze() if len(reconstructed_image.shape) == 2: - sitk_img = sitk.GetImageFromArray(reconstructed_image) + sitk_img = sitk.GetImageFromArray(reconstructed_image.T) # sitk_mask = sitk.GetImageFromArray(mask.astype(np.uint8).T) corrector = sitk.N4BiasFieldCorrectionImageFilter() @@ -405,9 +415,7 @@ def run_all(config) -> None: bias_n4 > 0, bias_n4, 1.0 ) imwrite( - reconstructed_image_filename.replace( - "corrected.tiff", "corrected_N4.tiff" - ), + reconstructed_image_filename.replace(".tiff", "_N4.tiff"), reconstructed_image_n4, ) diff --git a/mri_recon/distortions/__init__.py b/mri_recon/distortions/__init__.py index 874b248..24d63ca 100644 --- a/mri_recon/distortions/__init__.py +++ b/mri_recon/distortions/__init__.py @@ -2,6 +2,8 @@ BaseDistortion, DistortedKspaceMultiCoilMRI, SelfAdjointMultiplicativeMaskDistortion, + image_to_shifted_kspace, + shifted_kspace_to_image, ) from .biasfield import ( GaussianKspaceBiasField, From 37ab44219ccfcb6eb27356977ddf23cfbc56cfdd Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Tue, 9 Jun 2026 16:08:38 +0000 Subject: [PATCH 14/22] add prostate to inference plot script, remove debugging prints --- examples/fastmri_inference_plot.py | 40 +++++++++++++++++++----------- mri_recon/utils/plot.py | 7 ------ 2 files changed, 25 insertions(+), 22 deletions(-) diff --git a/examples/fastmri_inference_plot.py b/examples/fastmri_inference_plot.py index d9f73ac..79d02f1 100644 --- a/examples/fastmri_inference_plot.py +++ b/examples/fastmri_inference_plot.py @@ -22,7 +22,6 @@ ) from mri_recon.reconstruction import ( ConjugateGradientReconstructor, - OASISSinglecoilUnetReconstructor, choose_reconstructor, uses_oasis_centered_path, compatible_dataset_with_reconstructor, @@ -30,7 +29,8 @@ ) from mri_recon.utils import ( OasisCenteredFFTPhysics, - OasisSliceDataset, + OasisCenterSliceFolderDataset, + FastMRIProstateDataset, fastmri_measurement_to_oasis_kspace, kspace_to_image, save_kspace_plot, @@ -101,13 +101,13 @@ "width_y_fraction": 0.35, "center_x_fraction": 0.15, "center_y_fraction": -0.1, - "edge_gain": 0.3, + "edge_gain": 0.05, } }, { "GaussianBiasField": { "width_fraction": 0.35, - "edge_gain": 0.4, + "edge_gain": 0.05, } }, # {"AnisotropicLP": { @@ -201,6 +201,16 @@ def prepare_measurement_sample( and "coil_maps" in sample_batch[2] else None ) + elif dataset_name == "fastmri_prostate": + # reference image, shape: (B, W, H): dtype float32 + x = sample_batch[0].to(run_device) + + # add zero imaginary channel: + # (B, H, W) -> (B, 2, H, W) + x = torch.stack([x, torch.zeros_like(x)], dim=1) + + # (B, 2, H, W) + y = image_to_shifted_kspace(x) if use_oasis_fft_path: y = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) @@ -274,9 +284,6 @@ def build_physics_pair( selected_algorithms = ALGORITHMS if args.algorithm == "" else [args.algorithm] selected_distortions = DISTORTIONS if args.distortion == "" else [args.distortion] - REPORT_DIR = Path("reports") / Path(args.dataset + "_inference_plot") - REPORT_DIR.mkdir(parents=True, exist_ok=True) - # skip non-compatible algorithm-dataset pairs selected_algorithms = [ algo_name @@ -284,18 +291,11 @@ def build_physics_pair( if compatible_dataset_with_reconstructor(args.dataset, algo_name) ] - # set up report dir - if args.dataset not in ["fastmri", "oasis", "fastmri_multicoil", "cmrxrecon"]: - raise NotImplementedError(f"Invalid dataset: {args.dataset}") - # set up device, dataset, metrics device = dinv.utils.get_device() if args.dataset == "oasis": - split_csv = OASISSinglecoilUnetReconstructor.resolve_default_split_csv() - dataset = OasisSliceDataset( + dataset = OasisCenterSliceFolderDataset( data_path=args.source, - split_csv=split_csv, - sample_rate=0.6, ) elif args.dataset == "fastmri": dataset = dinv.datasets.FastMRISliceDataset(str(args.source), slice_index="middle") @@ -312,8 +312,18 @@ def build_physics_pair( dataset = dinv.datasets.CMRxReconSliceDataset( str(args.source), data_dir="SingleCoil/Cine/TrainingSet/FullSample", apply_mask=False ) + elif args.dataset == "fastmri_prostate": + dataset = FastMRIProstateDataset( + data_path=str(args.source), num_samples=args.num_samples, slice_index="middle" + ) + else: raise NotImplementedError(f"Invalid dataset: {args.dataset}") + + # set up report dir + REPORT_DIR = Path("reports") / Path(args.dataset + "_inference_plot") + REPORT_DIR.mkdir(parents=True, exist_ok=True) + metrics = [choose_metric(m) for m in METRICS] for i, batch in enumerate(iter(torch.utils.data.DataLoader(dataset))): diff --git a/mri_recon/utils/plot.py b/mri_recon/utils/plot.py index 35921b7..12b781c 100644 --- a/mri_recon/utils/plot.py +++ b/mri_recon/utils/plot.py @@ -53,18 +53,11 @@ def save_kspace_plot( ) -> None: """Save side-by-side log-magnitude visualizations of clean and distorted k-space.""" - print("transforming k-space to log-magnitude images for visualization...") - print(f"\tclean k-space shape: {clean_kspace.shape}") - print(f"\tdistorted k-space shape: {distorted_kspace.shape}") - images = [ ("Original k-space", _kspace_to_log_magnitude(clean_kspace)), ("Distorted k-space", _kspace_to_log_magnitude(distorted_kspace)), ] - print(f"clean k-space magnitude shape: {images[0][1].shape}") - print(f"distorted k-space magnitude shape: {images[1][1].shape}") - fig, axes = plt.subplots(1, 2, figsize=(8, 4), constrained_layout=True) fig.suptitle(f"Distortion: {distortion_label}") for ax, (title, image) in zip(axes, images, strict=True): From 059baa887f63c96ca7a4d80fc92b54d742cc350a Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Thu, 2 Jul 2026 09:41:20 +0000 Subject: [PATCH 15/22] add resolution distortions as separate loop to run_all, check, plot and export scripts --- examples/check_and_plot_results.py | 412 ++++++++++++++++++++++++++++ examples/export_examples.py | 342 +++++++++++++++++++++++ examples/run_all.py | 127 ++++++++- examples/test_dist_params.py | 76 +++++ mri_recon/distortions/__init__.py | 1 + mri_recon/distortions/resolution.py | 55 ++++ mri_recon/utils/prostate_adaptor.py | 19 +- 7 files changed, 1019 insertions(+), 13 deletions(-) create mode 100644 examples/check_and_plot_results.py create mode 100644 examples/export_examples.py create mode 100644 examples/test_dist_params.py diff --git a/examples/check_and_plot_results.py b/examples/check_and_plot_results.py new file mode 100644 index 0000000..fefb4fc --- /dev/null +++ b/examples/check_and_plot_results.py @@ -0,0 +1,412 @@ +import glob +import os + +import numpy as np +from tifffile import imread +import matplotlib.pyplot as plt + +result_folder = "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1_20260616" + +result_file_names = glob.glob(os.path.join(result_folder, "*.tiff")) + + +distortion_names = [ + "BaseDistortion", + "CartesianUndersamplingVariableDensity", + "CartesianUndersamplingUniformRandom", + "CartesianUndersamplingUniformRandomZeroACS", + "CartesianUndersamplingEquispaced", + "CartesianUndersamplingEquispacedZeroACS", + "PartialFourier", + "PhaseEncodeGhosting", + "SegmentedTranslationMotion", + "SegmentedRotationalMotion", + "TranslationMotion", + "RotationalMotion", + "OffCenterAnisotropicGaussianBiasField", + "GaussianBiasField", + "AnisotropicLP", + "HannTaperLP", + "KaiserTaperLP", + "GaussianNoise", + "IsotropicLP", + "RadialHighPassEmphasis", +] + +reconstruction_names = [ + "zero-filled", + "conjugate-gradient", + "ram", + "dip", + "tv-pgd", + "wavelet-fista", + "tv-fista", + "tv-pdhg", + "unet-fastmri", + "unet-oasis-acceleration4", + "unet-oasis-acceleration8", + "unet-oasis-acceleration10", +] + +samples = { + "fastmri_knee": ["1000000"], + "oasis": ["OAS1-0088-MR1"], + "fastmri_brain": ["AXFLAIR-200-6002452"], + "cmrxrecon": ["P001-cine-lax"], + "fastmri_prostate": ["AXT2-007"], +} + + +##### +# Check results +##### +check_results = True +if check_results: + missing_files = [] + for dataset_name, sample_names in samples.items(): + for sample_name in sample_names: + for distortion_name in distortion_names: + for reconstruction_name in reconstruction_names: + result_file_names_corr = glob.glob( + os.path.join( + result_folder, + f"image_{dataset_name}_{sample_name}_{distortion_name}*_{reconstruction_name}_corrected.tiff", + ) + ) + if len(result_file_names_corr) == 0: + missing_files.append( + f"image_{dataset_name}_{sample_name}_{distortion_name}_{reconstruction_name}_corrected.tiff" + ) + # else: + # result_file_names.remove(result_file_names_corr[0]) + result_file_names_uncorr = glob.glob( + os.path.join( + result_folder, + f"image_{dataset_name}_{sample_name}_{distortion_name}*_{reconstruction_name}_uncorrected.tiff", + ) + ) + if len(result_file_names_uncorr) == 0: + missing_files.append( + f"image_{dataset_name}_{sample_name}_{distortion_name}_{reconstruction_name}_uncorrected.tiff" + ) + # else: + # result_file_names.remove(result_file_names_uncorr[0]) + reference_file_name = os.path.join( + result_folder, f"image_{dataset_name}_{sample_name}_reference.tiff" + ) + if not os.path.exists(reference_file_name): + missing_files.append(reference_file_name) + kspace_ref_filename = os.path.join( + result_folder, f"kspace_{dataset_name}_{sample_name}_kspace_reference.tiff" + ) + if not os.path.exists(kspace_ref_filename): + missing_files.append(kspace_ref_filename) + distorted_kspace_file_name = os.path.join( + result_folder, + f"kspace_{dataset_name}_{sample_name}_{distortion_name}_distorted.tiff", + ) + if not os.path.exists(distorted_kspace_file_name): + missing_files.append(distorted_kspace_file_name) + + # found for prostate: + + found_prostate_files = sorted(glob.glob(os.path.join(result_folder, "*prostate*.tiff"))) + print(f"Found prostate files: {len(found_prostate_files)}") + + for file_name in found_prostate_files: + print(f"Found file: {file_name}") + + found_cmrxrecon_files = sorted(glob.glob(os.path.join(result_folder, "*cmrxrecon*.tiff"))) + print(f"Found cmrxrecon files: {len(found_cmrxrecon_files)}") + + for file_name in found_cmrxrecon_files: + print(f"Found file: {file_name}") + + for file_name in missing_files: + print("missing: ", file_name) + + +###### +# Create Plots for certain groups +###### + + +for dataset_name, sample_names in samples.items(): + for sample_name in sample_names: + print("looping over ", dataset_name, " ", sample_name) + + for part in range(2): + print("part ", part + 1) + + # always use CG in first column: + reconstruction_names = [ + recon for recon in reconstruction_names if recon != "conjugate-gradient" + ] + reconstruction_names_part = reconstruction_names[ + part * len(reconstruction_names) // 2 : (part + 1) * len(reconstruction_names) // 2 + ] + reconstruction_names_part = ["conjugate-gradient"] + reconstruction_names_part + + nr_rows = len(distortion_names) + nr_cols = len(reconstruction_names_part) * 2 + + fig, axes = plt.subplots( + int(nr_rows), int(nr_cols), figsize=(3 * nr_cols, 3 * nr_rows), squeeze=False + ) + fig.suptitle(f"{dataset_name} {sample_name} - part {part + 1}") + + # first row contains reference, BaseDistortion, without corrections + + axes[0, 0].set_title(f"{dataset_name} reference") + reference_file_name = os.path.join( + result_folder, f"image_{dataset_name}_{sample_name}_reference.tiff" + ) + if os.path.exists(reference_file_name): + img = imread(reference_file_name).squeeze() + if len(img.shape) == 3: + print(f"Warning: image has 3 dimensions: {img.shape}") + print(reference_file_name) + img = img[0, ...] + axes[0, 0].imshow(img, cmap="gray") + axes[0, 0].set_title(f"{dataset_name} reference") + axes[0, 0].xaxis.set_visible(False) + axes[0, 0].set_yticks([]) + axes[0, 0].set_ylabel("BaseDistortion", fontsize=12) + + else: + axes[0, 0].text( + 0.5, + 0.5, + "MISSING", + transform=axes[0, 0].transAxes, + fontsize=12, + color="red", + ha="center", + ) + axes[0, 0].axis("off") + + print("\tBaseDistortion") + for r_idx, reconstruction in enumerate(reconstruction_names_part): + print("\t\t", reconstruction) + if r_idx != 0: + result_file_names_uncorr = glob.glob( + os.path.join( + result_folder, + f"image_{dataset_name}_{sample_name}_BaseDistortion*_{reconstruction}_uncorrected.tiff", + ) + ) + if len(result_file_names_uncorr) > 0: + img = imread(result_file_names_uncorr[0]).squeeze() + if len(img.shape) == 3: + print(f"Warning: image has 3 dimensions: {img.shape}") + print(result_file_names_uncorr[0]) + img = img[0, ...] + axes[0, 2 * r_idx].imshow(img, cmap="gray") + axes[0, 2 * r_idx].set_title(f"{reconstruction} (u)") + axes[0, 2 * r_idx].axis("off") + else: + axes[0, 2 * r_idx].set_title(f"{reconstruction} (u)") + axes[0, 2 * r_idx].text( + 0.5, + 0.5, + "MISSING", + transform=axes[0, 2 * r_idx].transAxes, + fontsize=12, + color="red", + ha="center", + ) + axes[0, 2 * r_idx].axis("off") + + else: + axes[0, r_idx].axis("on") + axes[0, r_idx].xaxis.set_visible(False) + axes[0, r_idx].set_yticks([]) + axes[0, r_idx].set_ylabel("BaseDistortion", fontsize=12) + + result_file_names_corr = glob.glob( + os.path.join( + result_folder, + f"image_{dataset_name}_{sample_name}_BaseDistortion*_{reconstruction}_corrected.tiff", + ) + ) + if len(result_file_names_corr) > 0: + img = imread(result_file_names_corr[0]).squeeze() + if len(img.shape) == 3: + print(f"Warning: image has 3 dimensions: {img.shape}") + print(result_file_names_corr[0]) + img = img[0, ...] + axes[0, 2 * r_idx + 1].imshow(img, cmap="gray") + axes[0, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") + axes[0, 2 * r_idx + 1].axis("off") + else: + axes[0, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") + axes[0, 2 * r_idx + 1].text( + 0.5, + 0.5, + "MISSING", + transform=axes[0, 2 * r_idx + 1].transAxes, + fontsize=12, + color="red", + ha="center", + ) + axes[0, 2 * r_idx + 1].axis("off") + + # reconstruction methods in columns, distortions in rows + for d_idx, distortion in enumerate( + [dist for dist in distortion_names if dist != "BaseDistortion"] + ): + print("\t", distortion) + for r_idx, reconstruction in enumerate(reconstruction_names_part): + print("\t\t", reconstruction) + result_file_names_corr = glob.glob( + os.path.join( + result_folder, + f"image_{dataset_name}_{sample_name}_{distortion}*_{reconstruction}_corrected.tiff", + ) + ) + if len(result_file_names_corr) > 0: + img = imread(result_file_names_corr[0]).squeeze() + if len(img.shape) == 3: + print("Warning: image has 3 dimensions: [img.shape]") + print(result_file_names_corr[0]) + img = img[0, ...] + axes[d_idx + 1, 2 * r_idx + 1].imshow(img, cmap="gray") + axes[d_idx + 1, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") + axes[d_idx + 1, 2 * r_idx + 1].axis("off") + else: + axes[d_idx + 1, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") + axes[d_idx + 1, 2 * r_idx + 1].text( + 0.5, + 0.5, + "MISSING", + transform=axes[d_idx + 1, 2 * r_idx].transAxes, + fontsize=12, + color="red", + ha="center", + ) + axes[d_idx + 1, 2 * r_idx + 1].axis("off") + + result_file_names_uncorr = glob.glob( + os.path.join( + result_folder, + f"image_{dataset_name}_{sample_name}_{distortion}*_{reconstruction}_uncorrected.tiff", + ) + ) + if len(result_file_names_uncorr) > 0: + img_u = imread(result_file_names_uncorr[0]).squeeze() + if len(result_file_names_corr) > 0: + if (img == img_u).all(): + print( + f"Warning: corrected and uncorrected images are the same for {dataset_name}, {distortion}, {reconstruction}" + ) + if len(img_u.shape) == 3: + print(f"Warning: image has 3 dimensions: {img_u.shape}") + print(result_file_names_uncorr[0]) + img_u = img_u[0, ...] + axes[d_idx + 1, 2 * r_idx].imshow(img_u, cmap="gray") + axes[d_idx + 1, 2 * r_idx].set_title(f"{reconstruction} (u)") + axes[d_idx + 1, 2 * r_idx].axis("off") + else: + axes[d_idx + 1, 2 * r_idx].set_title(f"{reconstruction} (u)") + axes[d_idx + 1, 2 * r_idx].text( + 0.5, + 0.5, + "MISSING", + transform=axes[d_idx + 1, 2 * r_idx + 1].transAxes, + fontsize=12, + color="red", + ha="center", + ) + axes[d_idx + 1, 2 * r_idx].axis("off") + + if r_idx == 0: + axes[d_idx + 1, 0].axis("on") + axes[d_idx + 1, 0].xaxis.set_visible(False) + axes[d_idx + 1, 0].set_yticks([]) + if len(distortion) > 30: + # find the latest capital letter between second and 30th character and split there + for c in distortion[1:30]: + if c.isupper(): + cap_letter_idx = distortion[1:30].find(c) + if cap_letter_idx != -1: + str_distortion = ( + distortion[: cap_letter_idx + 1] + + "\n" + + distortion[cap_letter_idx + 1 :] + ) + else: + str_distortion = distortion[:30] + "\n" + distortion[30:50] + else: + str_distortion = distortion + axes[d_idx + 1, 0].set_ylabel(str_distortion, fontsize=12) + + plt.tight_layout() + plt.savefig( + os.path.join( + "/home/melanie.dohmen/ArtifactLab/reports", + f"summary_20260616_{dataset_name}_{sample_name}_part_{part + 1}.png", + ) + ) + + +###### +# Create Plots for each distortion and sample +###### + +for dataset_name, sample_names in samples.items(): + for sample_name in sample_names: + print("looping over ", dataset_name, " ", sample_name) + + for d_idx, distortion in enumerate( + [dist for dist in distortion_names if dist != "BaseDistortion"] + ): + print("\t", distortion) + + results_for_distortion = glob.glob( + os.path.join( + result_folder, f"image_{dataset_name}_{sample_name}_{distortion}*.tiff" + ) + ) + nr_rows = np.ceil(np.sqrt(len(results_for_distortion))) + nr_cols = np.ceil(len(results_for_distortion) / nr_rows) + + fig, axes = plt.subplots( + int(nr_rows), int(nr_cols), figsize=(3 * nr_cols, 3 * nr_rows), squeeze=False + ) + fig.suptitle(f"{dataset_name} {sample_name}") + + # sort results: + results_for_distortion_sorted = sorted( + results_for_distortion, key=lambda x: (x.split("_")[-2], x.split("_")[-1]) + ) + + # set BaseDistortion and CG first: + results_for_distortion_sorted = sorted( + results_for_distortion_sorted, + key=lambda x: ( + x.split("_")[-2] != "BaseDistortion", + x.split("_")[-2] != "conjugate-gradient", + ), + ) + + for r_idx, result_file_name in enumerate(results_for_distortion_sorted): + # split filename to get reconstruction name and corrected/uncorrected + reconstruction = result_file_name.split("_")[-2] + corrected = result_file_name.split("_")[-1].split(".")[0] + + img = imread(result_file_name).squeeze() + if len(img.shape) == 3: + print(f"Warning: image has 3 dimensions: {img.shape}") + print(result_file_name) + img = img[0, ...] + axes[r_idx // nr_cols, r_idx % nr_cols].imshow(img, cmap="gray") + axes[r_idx // nr_cols, r_idx % nr_cols].set_title(f"{reconstruction} (c)") + axes[r_idx // nr_cols, r_idx % nr_cols].axis("off") + + plt.tight_layout() + plt.savefig( + os.path.join( + "/home/melanie.dohmen/ArtifactLab/reports", + f"summary_20260616_{dataset_name}_{sample_name}_{distortion}.png", + ) + ) diff --git a/examples/export_examples.py b/examples/export_examples.py new file mode 100644 index 0000000..1bfa527 --- /dev/null +++ b/examples/export_examples.py @@ -0,0 +1,342 @@ +import os +import pandas as pd + +from tifffile import imread, imwrite + +# Structure for exporting examples: + + +# |-- property_01_pixel_resolution/ +# | |-- example_001 +# | |-- example_002/ +# | | |-- reference.npy +# | | |-- reference.png +# | | |-- degraded_1.npy +# | | |-- degraded_1.png +# | | |-- degraded_2.npy +# | | |-- degraded_2.png +# | | |-- degraded_3.npy +# | | |-- degraded_3.png +# | | `-- metadata.json <- Information about image sources and degration levels applied +# | |-- example_003 +# | |-- ... +# | `-- example_014 +# |-- property_02_texture_structure +# |-- property_03_image_contrast +# |-- ... +# |-- property_07_morphological_correctness `-- metadata.csv <- Spreadsheet with metadata from all examples from all properties +# ├── property_01_pixel_resolution +# │ ├── example_001 +# │ ├── example_002 +# │ │ ├── reference.npy +# │ │ ├── reference.png +# │ │ ├── degraded_1.npy +# │ │ ├── degraded_1.png +# │ │ ├── degraded_2.npy +# │ │ ├── degraded_2.png +# │ │ ├── degraded_3.npy +# │ │ ├── degraded_3.png +# │ │ └── metadata.json <- Information about image sources and degration levels applied +# │ ├── example_003 +# │ ├── ... +# │ └── example_014 +# ├── property_02_texture_structure +# ├── property_03_image_contrast +# ├── ... +# ├── property_07_morphological_correctness +# └── metadata.csv <- Spreadsheet with metadata from all examples from all properties +# . + + +# selected examples for each property: +properties = { + "property_01_pixel_resolution": { + "example_001": { + "reference": "image_fastmri_knee_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + "example_002": { + "reference": "image_fastmri_brain_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + }, + "property_02_sharpness": { + "example_001": { + "reference": "image_fastmri_knee_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + "example_002": { + "reference": "image_fastmri_brain_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + }, + "property_03_intensity_uniformity": { + "example_001": { + "reference": "image_fastmri_knee_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + "example_002": { + "reference": "image_fastmri_brain_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + }, + "property_04_noise_level": { + "example_001": { + "reference": "image_fastmri_knee_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + "example_002": { + "reference": "image_fastmri_brain_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + }, + "property_05_roi_homogeneity": { + "example_001": { + "reference": "image_fastmri_knee_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + "example_002": { + "reference": "image_fastmri_brain_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + }, + "property_06_local_signal_preservation": { + "example_001": { + "reference": "image_fastmri_knee_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + "example_002": { + "reference": "image_fastmri_brain_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + }, + "property_07_edges": { + "example_001": { + "reference": "image_fastmri_knee_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + "example_002": { + "reference": "image_fastmri_brain_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + }, + "property_08_contrast_preservation_of_anatomical_structures": { + "example_001": { + "reference": "image_fastmri_knee_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_knee_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + "example_002": { + "reference": "image_fastmri_brain_sample_0_reference.tiff", + "degraded": [ + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_uncorrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_zero-filled_corrected.tiff", + "image_fastmri_brain_sample_0_PartialFouriers=high_tv-pgd_corrected.tiff", + ], + }, + }, +} + +SOURCE_INFO = { + "fastmri_knee": { + "source_name": "fastMRI Knee", + "source_url": "https://fastmri.med.nyu.edu/", + "source_license": "internal research and educational purposes only", + }, + "fastmri_brain": { + "source_name": "fastMRI Brain", + "source_url": "https://fastmri.med.nyu.edu/", + "source_license": "internal research and educational purposes only", + }, + "fastmri_prostate": { + "source_name": "fastMRI Prostate", + "source_url": "https://fastmri.med.nyu.edu/", + "source_license": "internal research and educational purposes only", + }, + "oasis": { + "source_name": "OASIS", + "source_url": "https://sites.wustl.edu/oasisbrains/", + "source_license": "academic research purposes only", + }, + "cmrxrecon": { + "source_name": "CMRxRecon", + "source_url": "https://www.cmrxrecon.org/", + "source_license": "CC BY-NC-SA 4.0", + }, +} + + +def get_metadata(example_path: str) -> dict: + # remove parent directories from filename + example_path = os.path.basename(example_path) + example_path = example_path.replace(".tiff", "") + if "_N4" in example_path: + example_path = example_path.replace("_N4", "") + if "uncorrected" in example_path: + correction = "uncorrected" + example_path = example_path.replace("_uncorrected", "") + elif "corrected" in example_path: + correction = "corrected" + example_path = example_path.replace("_corrected", "") + else: + correction = "unknown" + # if no sample name is given, remove extra underscore + if "sample_0_" in example_path: + example_path = example_path.replace("sample_0_", "sample0_") + + filename_parts = example_path.split("_") + # [image_or_kspace, dataset_part1, (dataset_part2,) sample_name, distortion_or_reference, (reconstruction)] + if "reference" in filename_parts: + if len(filename_parts) == 4: + return { + **SOURCE_INFO[filename_parts[1]], # dataset_part1 + "sample_name": filename_parts[2], # sample_name + "distortion_type": "reference", + "correction": correction, + "reconstruction_method": "reference", + } + + elif len(filename_parts) == 5: + return { + **SOURCE_INFO[ + f"{filename_parts[1]}_{filename_parts[2]}" + ], # dataset_part1_dataset_part2 + "sample_name": filename_parts[3], # sample_name + "distortion_type": "reference", + "correction": correction, + "reconstruction_method": "reference", + } + else: + print(f"Warning: Unexpected filename format for reference example path: {example_path}") + elif len(filename_parts) == 5: + return { + **SOURCE_INFO[filename_parts[1]], # dataset_part1 + "sample_name": filename_parts[2], # sample_name + "distortion_type": filename_parts[3].split("=")[0], # distortion without parameters + "correction": correction, + "reconstruction_method": filename_parts[4], # reconstruction + } + elif len(filename_parts) == 6: + return { + **SOURCE_INFO[ + f"{filename_parts[1]}_{filename_parts[2]}" + ], # dataset_part1_dataset_part2 + "sample_name": filename_parts[3], # sample_name + "distortion_type": filename_parts[4].split("=")[0], # distortion without parameters + "correction": correction, + "reconstruction_method": filename_parts[5], # reconstruction + } + else: + print(f"Warning: Unexpected filename format for example path: {example_path}") + + return { + "source": "unknown", + "sample_name": "unknown", + "distortion_type": "unknown", + "correction": correction, + "reconstruction_method": "unknown", + } + + +result_path = "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1_20260616/" +export_path = "/home/melanie.dohmen/ArtifactLab/reports/exported_examples/" +os.makedirs(result_path, exist_ok=True) + + +metadata_df_list = [] + +for property_name, property_examples in properties.items(): + property_path = os.path.join(export_path, property_name) + os.makedirs(property_path, exist_ok=True) + + for example_name, example_data in property_examples.items(): + example_path = os.path.join(property_path, example_name) + os.makedirs(example_path, exist_ok=True) + + # Save reference image + old_reference_image_path = os.path.join(result_path, example_data["reference"]) + new_reference_image_path = os.path.join(example_path, "reference.tiff") + + imwrite(new_reference_image_path, imread(old_reference_image_path)) + # Code to save the reference image using example_data["reference"] + + # Save degraded images + for i, degraded_image in enumerate(example_data["degraded"]): + old_degraded_image_path = os.path.join(result_path, degraded_image) + new_degraded_image_path = os.path.join(example_path, f"degraded_{i + 1}.tiff") + imwrite(new_degraded_image_path, imread(old_degraded_image_path)) + + metadata = { + "property": property_name, + "example": example_name, + "degraded_image_index": i + 1, + "relative_path": os.path.relpath(new_degraded_image_path, result_path), + **get_metadata( + old_degraded_image_path + ), # Function to extract metadata info from filename + } + + metadata_df_list.append(metadata) + +metadata_df = pd.DataFrame(metadata_df_list) +metadata_csv_path = os.path.join(export_path, "metadata.csv") +metadata_df.to_csv(metadata_csv_path, index=False) diff --git a/examples/run_all.py b/examples/run_all.py index e58833f..525876e 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -11,7 +11,6 @@ import numpy as np import tqdm - sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from datetime import datetime @@ -23,6 +22,7 @@ from mri_recon.distortions import ( BaseDistortion, DistortedKspaceMultiCoilMRI, + ResolutionReductionByKspaceCropping, choose_distortion_with_params, image_to_shifted_kspace, ) @@ -55,6 +55,7 @@ def get_measurement_sample( """ coil_maps = None if dataset_name == "oasis": + sample_name = sample_batch["subject_id"][0].replace("_", "-") # reference image, shape: (B, 2, H, W) dtype: float32 x = sample_batch["x"].to(run_device) # centered k-space data, shape: (B, 2, H, W) dtype: float32 @@ -62,7 +63,10 @@ def get_measurement_sample( # k-space data, shape: (B, 2, H, W) dtype: float32 # y = oasis_kspace_to_fastmri_measurement(y_centered) y = image_to_shifted_kspace(x) + elif dataset_name == "fastmri_knee": + # sample name must be fetched from dataset directly + sample_name = None # reference image, shape: (B, 1, H/2, H/2) dtype: float32 x = sample_batch[0].to(run_device) # kspace data, shape: (B, 2, H, W) dtype: float32 @@ -72,10 +76,13 @@ def get_measurement_sample( # reconstructed reference image: # shape: (B, 1, H, W) dtype: float32 elif dataset_name == "fastmri_brain": + # sample name must be fetched from dataset directly + sample_name = None # reference image, shape: (B, 1, H/2, H/2) dtype: float32 x = sample_batch[0].to(run_device) # kspace data, shape: (B, 2, num_coils, H, W) dtype: float32 y = sample_batch[1].to(run_device) + # coil maps, shape: (B, num_coils, H, W) dtype: complex64 coil_maps = ( sample_batch[2]["coil_maps"].to(run_device) @@ -88,6 +95,8 @@ def get_measurement_sample( y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) elif dataset_name == "cmrxrecon": + # sample name must be fetched from dataset directly + sample_name = None # reference image, shape: (B, 2, n_timepoints, (n_coils), H, W) x = sample_batch[0].to(run_device) # k-space data, shape: (B, 2, n_timepoints, (n_coils), H, W) dtype: float32 @@ -113,7 +122,8 @@ def get_measurement_sample( elif dataset_name == "fastmri_prostate": # reference image, shape: (B, W, H): dtype float32 - x = sample_batch[0].to(run_device) + x = sample_batch["image"].to(run_device) + sample_name = sample_batch["sample_name"][0] # add zero imaginary channel: # (B, H, W) -> (B, 2, H, W) @@ -130,7 +140,7 @@ def get_measurement_sample( print("y_centered: ", y_centered.shape) print("coil_maps: ", coil_maps.shape if coil_maps is not None else "None") - return x, y, y_centered, coil_maps + return x, y, y_centered, coil_maps, sample_name def run_all(config) -> None: @@ -178,16 +188,38 @@ def run_all(config) -> None: break print(f"{dataset_name} sample {i}...") - x_reference, y, y_centered, coil_maps = get_measurement_sample( + x_reference, y, y_centered, coil_maps, sample_name = get_measurement_sample( sample_batch=batch, dataset_name=dataset_name, run_device=device, ) + if sample_name is None: + if dataset_name == "fastmri_knee" or dataset_name == "fastmri_brain": + fname, _, _ = dataset.samples[i] + sample_name = os.path.basename(fname).split(".")[0] + sample_name = ( + sample_name.replace("brain_", "") + .replace("knee_", "") + .replace("file_", "") + .replace("file", "") + .replace("_", "-") + ) + elif dataset_name == "cmrxrecon": + fname, _, _ = dataset.samples[i] + patient_id = os.path.basename(os.path.dirname(fname)) + sample_name = ( + f"{patient_id}-{os.path.basename(fname).split('.')[0].replace('_', '-')}" + ) + else: + sample_name = f"sample{i}" + + print("sample_name: ", sample_name) + # save reference image imwrite( os.path.join( - config["results_dir"], f"image_{dataset_name}_sample_{i}_reference.tiff" + config["results_dir"], f"image_{dataset_name}_{sample_name}_reference.tiff" ), convert_image_for_save(x_reference), ) @@ -240,14 +272,14 @@ def run_all(config) -> None: imwrite( os.path.join( config["results_dir"], - f"kspace_{dataset_name}_sample_{i}_reference.tiff", + f"kspace_{dataset_name}_{sample_name}_reference.tiff", ), _kspace_to_log_magnitude(y).numpy(), ) imwrite( os.path.join( config["results_dir"], - f"kspace_{dataset_name}_sample_{i}_{distortion_name_with_params}.tiff", + f"kspace_{dataset_name}_{sample_name}_{distortion_name_with_params}.tiff", ), _kspace_to_log_magnitude(y_distorted).numpy(), ) @@ -272,14 +304,14 @@ def run_all(config) -> None: imwrite( os.path.join( config["results_dir"], - f"image_{dataset_name}_sample_{i}_{distortion_name_with_params}_{reconstructor_name}_uncorrected.tiff", + f"image_{dataset_name}_{sample_name}_{distortion_name_with_params}_{reconstructor_name}_uncorrected.tiff", ), convert_image_for_save(x_uncorrected), ) imwrite( os.path.join( config["results_dir"], - f"image_{dataset_name}_sample_{i}_{distortion_name_with_params}_{reconstructor_name}_corrected.tiff", + f"image_{dataset_name}_{sample_name}_{distortion_name_with_params}_{reconstructor_name}_corrected.tiff", ), convert_image_for_save(x_corrected), ) @@ -426,6 +458,83 @@ def run_all(config) -> None: pbar.update(1) + # apply true resolution reduction (image matrix size) by k-space cropping + for factor in config["resolution_reduction_factors"]: + print(f"Applying resolution reduction with factor {factor}") + + kspace_crop = ResolutionReductionByKspaceCropping( + crop_fraction=1.0 / factor, img_size=x_reference.shape[-2:] + ) + y_distorted = kspace_crop._apply_crop(y) + + if coil_maps is not None: + coil_maps_channels = torch.view_as_real(coil_maps) + coil_maps_channels_lowres_realnn = torch.nn.functional.interpolate( + coil_maps_channels[..., 0], scale_factor=0.5, mode="nearest" + ) + coil_maps_channels_lowres_imagnn = torch.nn.functional.interpolate( + coil_maps_channels[..., 1], scale_factor=0.5, mode="nearest" + ) + coil_maps_lowresnn = torch.view_as_complex( + torch.stack( + [coil_maps_channels_lowres_realnn, coil_maps_channels_lowres_imagnn], + dim=-1, + ) + ) + else: + coil_maps_lowresnn = None + + physics_distorted = DistortedKspaceMultiCoilMRI( + BaseDistortion(), + img_size=(int(y.shape[-2] / factor), int(y.shape[-1] / factor)), + coil_maps=coil_maps_lowresnn, + device=device, + ) + + for reconstructor_name in config["reconstruction_algorithms"]: + # only run on reconstructors, that use the fastmri-like k-space + if not uses_oasis_centered_path(reconstructor_name): + print(f"\t\t{reconstructor_name} ...") + start = datetime.now() + if compatible_dataset_with_reconstructor(dataset_name, reconstructor_name): + reconstructor = choose_reconstructor( + reconstructor_name, + img_size=y_distorted.shape[-2:], + device=device, + verbose=config["verbose"], + ).to(device) + + imwrite( + os.path.join( + config["results_dir"], + f"kspace_{dataset_name}_{sample_name}_ReduceResolutionf={factor}.tiff", + ), + _kspace_to_log_magnitude(y_distorted).numpy(), + ) + + # actual reconstruction with the selected reconstructor + try: + x_corrected = reconstructor(y_distorted, physics_distorted) + + # restore original image size from reconstructed image by upsampling (and cropping if necessary) + x_corrected = kspace_crop._upsample_back(x_corrected) + + # save reconstructed images + imwrite( + os.path.join( + config["results_dir"], + f"image_{dataset_name}_{sample_name}_ReduceResolutionf={factor}_{reconstructor_name}_corrected.tiff", + ), + convert_image_for_save(x_corrected), + ) + + print(f"\t\t... done in {datetime.now() - start}") + + except Exception as e: + print(f"Error using {reconstructor_name}: {e}") + else: + print(f"\t\t ... not compatible with {dataset_name}") + if __name__ == "__main__": # read config file in yaml format as first argument from commmand line diff --git a/examples/test_dist_params.py b/examples/test_dist_params.py new file mode 100644 index 0000000..5f3a354 --- /dev/null +++ b/examples/test_dist_params.py @@ -0,0 +1,76 @@ +import os +import sys + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +from run_all import run_all + +if __name__ == "__main__": + results_dir = "/home/melanie.dohmen/ArtifactLab/reports/" + + # create config: + config = { + "data": { + "fastmri_knee": "/home/melanie.dohmen/ArtifactLab/data/singlecoil_val", + "oasis": "/home/melanie.dohmen/ArtifactLab/data/oasis", + "fastmri_brain": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_multicoil_brain_test", + "cmrxrecon": "/home/melanie.dohmen/ArtifactLab/data/CMRxRecon", + "fastmri_prostate": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_prostate_T2_IDS_001_020", + }, + "distortions": [ + { + "CartesianUndersamplingEquispacedZeroACS": {"keep_fraction": 0.25}, + }, + ], + "reconstruction_algorithms": [ + "zero-filled", + "conjugate-gradient", + ], + "add_N4Correction": False, + "num_samples": 1, + "verbose": True, + "results_dir": results_dir, + } + + distortions = [ + # {"IsotropicLP": { + # "radius_fraction": [0.01, 0.05, 0.1, 0.15, 0.2, 0.5] + # }}, + # {"GaussianKspaceBiasField": { + # "width_fraction": [0.1, 0.35, 0.5, 0.35, 0.35, 0.35], + # "edge_gain": [ 0.4, 0.4, 0.4, 0.1, 0.2, 0.6, 0.8], + # }}, + # 0.01 width fraction toooo small! + # {"GaussianBiasField": { + # "width_fraction": [0.1, 0.15, 0.2, 0.5, 0.1, 0.15, 0.2, 0.5, 0.1, 0.15, 0.2, 0.5,], + # "edge_gain": [ 0.05, 0.05, 0.05, 0.05, 0.1, 0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.2,], + # }}, + # {"OffCenterAnisotropicGaussianBiasField": { + # "width_x_fraction": [ 0.1, 0.15, 0.2, 0.35, 0.1, 0.15, 0.2, 0.35, 0.1, 0.15, 0.2, 0.35], + # "width_y_fraction": [0.15, 0.2, 0.35, 0.1, 0.2, 0.35, 0.1, 0.15, 0.35, 0.1, 0.15, 0.2, ], + # "center_x_fraction": [0.15, 0.15, 0.15,0.15, 0.15,0.15,0.15,0.15, 0.15,0.15, 0.15, 0.15,], + # "center_y_fraction": [-0.1, -0.1, -0.1, -0.1, -0.1,-0.1,-0.1,-0.1, -0.1,-0.1,-0.1,-0.1,], + # "edge_gain": [0.05,0.05,0.05,0.05, 0.1, 0.1, 0.1, 0.1, 0.5, 0.5, 0.5, 0.5, ], + # }}, + { + "CartesianUndersamplingEquispacedZeroACS": { + "keep_fraction": [0.1, 0.25, 0.5, 0.75, 0.85, 0.9, 0.95, 0.98] + }, + }, + ] + + for d_idx, distortion_dict in enumerate(distortions): + for distortion_name, dist_params in distortion_dict.items(): + nr_param_values = len(dist_params[list(dist_params.keys())[0]]) + for v_idx in range(nr_param_values): + single_value_distortion_dict = { + distortion_name: { + param: param_values[v_idx] for param, param_values in dist_params.items() + } + } + config["distortions"].append(single_value_distortion_dict) + + config["results_dir"] = os.path.join(results_dir, f"test_params_{distortion_name}") + + run_all(config) diff --git a/mri_recon/distortions/__init__.py b/mri_recon/distortions/__init__.py index 24d63ca..d401f08 100644 --- a/mri_recon/distortions/__init__.py +++ b/mri_recon/distortions/__init__.py @@ -24,6 +24,7 @@ IsotropicResolutionReduction, KaiserTaperResolutionReduction, RadialHighPassEmphasisDistortion, + ResolutionReductionByKspaceCropping, ) from .undersampling import CartesianUndersampling, PartialFourierDistortion from .utils import choose_distortion_with_params diff --git a/mri_recon/distortions/resolution.py b/mri_recon/distortions/resolution.py index 399a037..b882f3e 100644 --- a/mri_recon/distortions/resolution.py +++ b/mri_recon/distortions/resolution.py @@ -351,3 +351,58 @@ def _mask(self, shape: tuple[int, ...], device: torch.device) -> torch.Tensor: transition = transition.clamp(0.0, 1.0) transition = transition * transition * (3.0 - 2.0 * transition) return 1.0 + self.alpha * transition + + +class ResolutionReductionByKspaceCropping: + """ + Reduces the resolution of an image by cropping its k-space representation. + The cropping is done by keeping only a fraction of the k-space data, specified by the `crop_fraction` parameter. + + The image with lower resolution can be upsampled back to the original size by the `_upsample_back` method, but this will not recover the lost high-frequency information. + + This distortion is not compatible with BaseDistortion, because it changes the dimensions of the given k-space Tensor y. + """ + + def __init__(self, crop_fraction: float = 0.5, img_size: tuple[int, int] = None) -> None: + super().__init__() + if not 0.0 < crop_fraction <= 1.0: + raise ValueError("crop_fraction must be in (0, 1]") + self.crop_fraction = crop_fraction + self.img_size = img_size + + def _apply_crop(self, y: torch.Tensor) -> torch.Tensor: + if self.img_size is None: + self.img_size = y.shape[-2:] + # y is assumed to be in k-space + # crop k-space to reduce resolution + w = y.shape[-2] + h = y.shape[-1] + w_lowres = int(w * self.crop_fraction) + h_lowres = int(h * self.crop_fraction) + frac_diffw = (w - w_lowres) // 2 + frac_diffh = (h - h_lowres) // 2 + y_cropped = y[..., frac_diffw : w_lowres + frac_diffw, frac_diffh : h_lowres + frac_diffh] + return y_cropped + + def _upsample_back(self, y_cropped: torch.Tensor) -> torch.Tensor: + # upsample back to original size + + y_upsampled = torch.nn.functional.interpolate( + y_cropped, scale_factor=1 / self.crop_fraction, mode="nearest" + ) + + # center crop to original image size: + if self.img_size is not None: + if ( + y_upsampled.shape[-2] != self.img_size[0] + or y_upsampled.shape[-1] != self.img_size[1] + ): + frac_diffw = (y_upsampled.shape[-2] - self.img_size[0]) // 2 + frac_diffh = (y_upsampled.shape[-1] - self.img_size[1]) // 2 + y_upsampled = y_upsampled[ + ..., + frac_diffw : frac_diffw + self.img_size[0], + frac_diffh : frac_diffh + self.img_size[1], + ] + + return y_upsampled diff --git a/mri_recon/utils/prostate_adaptor.py b/mri_recon/utils/prostate_adaptor.py index c7b1d6a..a21184e 100644 --- a/mri_recon/utils/prostate_adaptor.py +++ b/mri_recon/utils/prostate_adaptor.py @@ -12,15 +12,16 @@ def __init__( self.data_path = data_path self.num_samples = num_samples self.slice_index = slice_index - self.image_data = self.get_image_data() + self.image_data, self.sample_names = self.get_image_data() if num_samples is not None: self.image_data = self.image_data[:num_samples] else: self.num_samples = len(self.image_data) - def get_image_data(self) -> np.ndarray: + def get_image_data(self) -> tuple[np.ndarray, list[str]]: image_result_list = [] + sample_name_list = [] for sample_idx, filename in enumerate(glob.glob(os.path.join(self.data_path, "*.h5"))): if (self.num_samples is not None) and (sample_idx >= self.num_samples): break @@ -38,12 +39,22 @@ def get_image_data(self) -> np.ndarray: image_result_list.extend( [image_recon[i, :, :] for i in range(image_recon.shape[0])] ) + sample_name = ( + os.path.basename(filename) + .split(".")[0] + .replace("file_prostate_", "") + .replace("_", "-") + ) + sample_name_list.append(sample_name) - return image_result_list + return image_result_list, sample_name_list def __len__(self) -> int: return len(self.image_data) def __getitem__(self, idx: int) -> torch.Tensor: # add batch dimension and convert to torch.Tensor - return torch.from_numpy(self.image_data[idx]).unsqueeze(0) + return { + "image": torch.from_numpy(self.image_data[idx]).unsqueeze(0), + "sample_name": self.sample_names[idx], + } From b9fac40d3fcf8da85238c1ab21f5ab130b4c01d5 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Thu, 2 Jul 2026 10:48:52 +0000 Subject: [PATCH 16/22] add simple itk dependency --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4a45111..36497c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,7 @@ dependencies = [ "pytest>=9.0.2", "python-certifi-win32>=1.6.1", "sigpy>=0.1.27", - "SimpleITK>=2.5.5", + "simpleitk==2.5.5", "torch>=2.11.0", "torchmetrics>=1.9.0", "torchvision>=0.26.0", diff --git a/uv.lock b/uv.lock index 1b5421b..ee48044 100644 --- a/uv.lock +++ b/uv.lock @@ -171,7 +171,7 @@ requires-dist = [ { name = "pytest", specifier = ">=9.0.2" }, { name = "python-certifi-win32", specifier = ">=1.6.1" }, { name = "sigpy", specifier = ">=0.1.27" }, - { name = "simpleitk", specifier = ">=2.5.5" }, + { name = "simpleitk", specifier = "==2.5.5" }, { name = "torch", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", specifier = ">=2.11.0" }, { name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.11.0", index = "https://download.pytorch.org/whl/cpu" }, { name = "torch", marker = "sys_platform == 'linux'", specifier = ">=2.11.0", index = "https://download.pytorch.org/whl/cu128" }, From 66ba6b97cdaea2400e0995bda1124c499bd7798e Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Wed, 8 Jul 2026 22:11:02 +0000 Subject: [PATCH 17/22] replace sample_0 by sample name --- examples/run_all.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/run_all.py b/examples/run_all.py index 525876e..2a68e5f 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -364,14 +364,14 @@ def run_all(config) -> None: imwrite( os.path.join( config["results_dir"], - f"kspace_centered_{dataset_name}_sample_{i}_reference.tiff", + f"kspace_centered_{dataset_name}_{sample_name}_reference.tiff", ), _kspace_to_log_magnitude(y_centered).numpy(), ) imwrite( os.path.join( config["results_dir"], - f"kspace_centered_{dataset_name}_sample_{i}_{distortion_name_with_params}.tiff", + f"kspace_centered_{dataset_name}_{sample_name}_{distortion_name_with_params}.tiff", ), _kspace_to_log_magnitude(y_distorted).numpy(), ) @@ -395,14 +395,14 @@ def run_all(config) -> None: imwrite( os.path.join( config["results_dir"], - f"image_{dataset_name}_sample_{i}_{distortion_name_with_params}_{reconstructor_name}_uncorrected.tiff", + f"image_{dataset_name}_{sample_name}_{distortion_name_with_params}_{reconstructor_name}_uncorrected.tiff", ), convert_image_for_save(x_uncorrected), ) imwrite( os.path.join( config["results_dir"], - f"image_{dataset_name}_sample_{i}_{distortion_name_with_params}_{reconstructor_name}_corrected.tiff", + f"image_{dataset_name}_{sample_name}_{distortion_name_with_params}_{reconstructor_name}_corrected.tiff", ), convert_image_for_save(x_corrected), ) @@ -410,7 +410,7 @@ def run_all(config) -> None: except Exception as e: print( - f"\t\tError using {reconstructor_name} with distortion {distortion_name_with_params} on sample {i}: {e}" + f"\t\tError using {reconstructor_name} with distortion {distortion_name_with_params} on {sample_name}: {e}" ) else: From ad7107e56276d5a2b9103ddcd667a52de528546b Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Wed, 8 Jul 2026 22:30:17 +0000 Subject: [PATCH 18/22] missing correct params --- examples/config.yaml | 36 +++++++++++++++++++++--------------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/examples/config.yaml b/examples/config.yaml index 304002b..4d2d3ae 100644 --- a/examples/config.yaml +++ b/examples/config.yaml @@ -8,37 +8,43 @@ data: distortions: - "BaseDistortion": {} - "CartesianUndersamplingVariableDensity": - "keep_fraction": "keep_fraction" - "center_fraction": "center_fraction" + "keep_fraction": 0.75 + "center_fraction": 0.125 - "CartesianUndersamplingUniformRandom": - "keep_fraction": "keep_fraction" - "center_fraction": "center_fraction" + "keep_fraction": 0.75 + "center_fraction": 0.125 - "CartesianUndersamplingUniformRandomZeroACS": - "keep_fraction": "keep_fraction" + "keep_fraction": 0.95 - "CartesianUndersamplingEquispaced": - "keep_fraction": "keep_fraction" - "center_fraction": "center_fraction" + "keep_fraction": 0.75 + "center_fraction": 0.125 - "CartesianUndersamplingEquispacedZeroACS": - "keep_fraction": "keep_fraction" + "keep_fraction": 0.95 - "PartialFourier": "side": "high" - "PhaseEncodeGhosting": "line_period": 2, "line_offset": 1, - "phase_error_radians": torch.pi / 2, + #"phase_error_radians": torch.pi / 2, + "phase_error_degrees": 90.0 "corrupted_line_scale": 1.0 - "SegmentedTranslationMotion": - "shift_x_pixels": [0.0, 20.0, 50.0, -50.0] - "shift_y_pixels": [0.0, 10.0, -20.0, 20.0] + #"shift_x_pixels": [0.0, 20.0, 50.0, -50.0] + #"shift_y_pixels": [0.0, 10.0, -20.0, 20.0] + "shift_x_pixels": [2.0, 0.0, -3.0, -5.0] + "shift_y_pixels": [-1.0, 1.0, 2.0, 3.0] - "SegmentedRotationalMotion": #"angle_radians": [0.0, torch.pi / 20, -torch.pi / 24, torch.pi / 16] - "angle_degrees": [0.0, 18.0, -15.0, 22.5] + "angle_degrees": [0.0, 2.0, -1.0, -3] - "TranslationMotion": - "shift_x_pixels": 60 - "shift_y_pixels": 10 + #"shift_x_pixels": 60 + #"shift_y_pixels": 10 + "shift_x_pixels": 5 + "shift_y_pixels": 2 - "RotationalMotion": # angle_radians=torch.pi / 6 - "angle_degrees": 60.0 + # angle_degrees: 60 + "angle_degrees": 2.0 # - "OffCenterAnisotropicGaussianKspaceBiasField": # "width_x_fraction": 0.2 # "width_y_fraction": 0.35 From affd9da29b1d4d0bea568fd64fc035f50d737f1e Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Wed, 8 Jul 2026 22:30:45 +0000 Subject: [PATCH 19/22] improve distortion plots --- examples/check_and_plot_results.py | 504 ++++++++++++++++------------- 1 file changed, 275 insertions(+), 229 deletions(-) diff --git a/examples/check_and_plot_results.py b/examples/check_and_plot_results.py index fefb4fc..1be91d2 100644 --- a/examples/check_and_plot_results.py +++ b/examples/check_and_plot_results.py @@ -5,7 +5,8 @@ from tifffile import imread import matplotlib.pyplot as plt -result_folder = "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1_20260616" +result_folder = "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1" +# result_folder = "/home/melanie.dohmen/ArtifactLab/reports/test_params_AnisotropicLP" result_file_names = glob.glob(os.path.join(result_folder, "*.tiff")) @@ -53,14 +54,14 @@ "oasis": ["OAS1-0088-MR1"], "fastmri_brain": ["AXFLAIR-200-6002452"], "cmrxrecon": ["P001-cine-lax"], - "fastmri_prostate": ["AXT2-007"], + "fastmri_prostate": ["AXT2-013", "AXT2-007"], } ##### # Check results ##### -check_results = True +check_results = False if check_results: missing_files = [] for dataset_name, sample_names in samples.items(): @@ -127,286 +128,331 @@ ###### -# Create Plots for certain groups +# Create Plots for each sample ###### +create_sample_summary = False -for dataset_name, sample_names in samples.items(): - for sample_name in sample_names: - print("looping over ", dataset_name, " ", sample_name) - - for part in range(2): - print("part ", part + 1) - - # always use CG in first column: - reconstruction_names = [ - recon for recon in reconstruction_names if recon != "conjugate-gradient" - ] - reconstruction_names_part = reconstruction_names[ - part * len(reconstruction_names) // 2 : (part + 1) * len(reconstruction_names) // 2 - ] - reconstruction_names_part = ["conjugate-gradient"] + reconstruction_names_part - - nr_rows = len(distortion_names) - nr_cols = len(reconstruction_names_part) * 2 - - fig, axes = plt.subplots( - int(nr_rows), int(nr_cols), figsize=(3 * nr_cols, 3 * nr_rows), squeeze=False - ) - fig.suptitle(f"{dataset_name} {sample_name} - part {part + 1}") - - # first row contains reference, BaseDistortion, without corrections - - axes[0, 0].set_title(f"{dataset_name} reference") - reference_file_name = os.path.join( - result_folder, f"image_{dataset_name}_{sample_name}_reference.tiff" - ) - if os.path.exists(reference_file_name): - img = imread(reference_file_name).squeeze() - if len(img.shape) == 3: - print(f"Warning: image has 3 dimensions: {img.shape}") - print(reference_file_name) - img = img[0, ...] - axes[0, 0].imshow(img, cmap="gray") - axes[0, 0].set_title(f"{dataset_name} reference") - axes[0, 0].xaxis.set_visible(False) - axes[0, 0].set_yticks([]) - axes[0, 0].set_ylabel("BaseDistortion", fontsize=12) - - else: - axes[0, 0].text( - 0.5, - 0.5, - "MISSING", - transform=axes[0, 0].transAxes, - fontsize=12, - color="red", - ha="center", +if create_sample_summary: + for dataset_name, sample_names in samples.items(): + for sample_name in sample_names: + print("looping over ", dataset_name, " ", sample_name) + + for part in range(2): + print("part ", part + 1) + + # always use CG in first column: + reconstruction_names = [ + recon for recon in reconstruction_names if recon != "conjugate-gradient" + ] + reconstruction_names_part = reconstruction_names[ + part * len(reconstruction_names) // 2 : (part + 1) + * len(reconstruction_names) + // 2 + ] + reconstruction_names_part = ["conjugate-gradient"] + reconstruction_names_part + + nr_rows = len(distortion_names) + nr_cols = len(reconstruction_names_part) * 2 + + fig, axes = plt.subplots( + int(nr_rows), int(nr_cols), figsize=(3 * nr_cols, 3 * nr_rows), squeeze=False ) - axes[0, 0].axis("off") + fig.suptitle(f"{dataset_name} {sample_name} - part {part + 1}") - print("\tBaseDistortion") - for r_idx, reconstruction in enumerate(reconstruction_names_part): - print("\t\t", reconstruction) - if r_idx != 0: - result_file_names_uncorr = glob.glob( - os.path.join( - result_folder, - f"image_{dataset_name}_{sample_name}_BaseDistortion*_{reconstruction}_uncorrected.tiff", - ) - ) - if len(result_file_names_uncorr) > 0: - img = imread(result_file_names_uncorr[0]).squeeze() - if len(img.shape) == 3: - print(f"Warning: image has 3 dimensions: {img.shape}") - print(result_file_names_uncorr[0]) - img = img[0, ...] - axes[0, 2 * r_idx].imshow(img, cmap="gray") - axes[0, 2 * r_idx].set_title(f"{reconstruction} (u)") - axes[0, 2 * r_idx].axis("off") - else: - axes[0, 2 * r_idx].set_title(f"{reconstruction} (u)") - axes[0, 2 * r_idx].text( - 0.5, - 0.5, - "MISSING", - transform=axes[0, 2 * r_idx].transAxes, - fontsize=12, - color="red", - ha="center", - ) - axes[0, 2 * r_idx].axis("off") - - else: - axes[0, r_idx].axis("on") - axes[0, r_idx].xaxis.set_visible(False) - axes[0, r_idx].set_yticks([]) - axes[0, r_idx].set_ylabel("BaseDistortion", fontsize=12) + # first row contains reference, BaseDistortion, without corrections - result_file_names_corr = glob.glob( - os.path.join( - result_folder, - f"image_{dataset_name}_{sample_name}_BaseDistortion*_{reconstruction}_corrected.tiff", - ) + axes[0, 0].set_title(f"{dataset_name} reference") + reference_file_name = os.path.join( + result_folder, f"image_{dataset_name}_{sample_name}_reference.tiff" ) - if len(result_file_names_corr) > 0: - img = imread(result_file_names_corr[0]).squeeze() + if os.path.exists(reference_file_name): + img = imread(reference_file_name).squeeze() if len(img.shape) == 3: print(f"Warning: image has 3 dimensions: {img.shape}") - print(result_file_names_corr[0]) + print(reference_file_name) img = img[0, ...] - axes[0, 2 * r_idx + 1].imshow(img, cmap="gray") - axes[0, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") - axes[0, 2 * r_idx + 1].axis("off") + axes[0, 0].imshow(img, cmap="gray") + axes[0, 0].set_title(f"{dataset_name} reference") + axes[0, 0].xaxis.set_visible(False) + axes[0, 0].set_yticks([]) + axes[0, 0].set_ylabel("BaseDistortion", fontsize=12) + else: - axes[0, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") - axes[0, 2 * r_idx + 1].text( + axes[0, 0].text( 0.5, 0.5, "MISSING", - transform=axes[0, 2 * r_idx + 1].transAxes, + transform=axes[0, 0].transAxes, fontsize=12, color="red", ha="center", ) - axes[0, 2 * r_idx + 1].axis("off") + axes[0, 0].axis("off") - # reconstruction methods in columns, distortions in rows - for d_idx, distortion in enumerate( - [dist for dist in distortion_names if dist != "BaseDistortion"] - ): - print("\t", distortion) + print("\tBaseDistortion") for r_idx, reconstruction in enumerate(reconstruction_names_part): print("\t\t", reconstruction) + if r_idx != 0: + result_file_names_uncorr = glob.glob( + os.path.join( + result_folder, + f"image_{dataset_name}_{sample_name}_BaseDistortion*_{reconstruction}_uncorrected.tiff", + ) + ) + if len(result_file_names_uncorr) > 0: + img = imread(result_file_names_uncorr[0]).squeeze() + if len(img.shape) == 3: + print(f"Warning: image has 3 dimensions: {img.shape}") + print(result_file_names_uncorr[0]) + img = img[0, ...] + axes[0, 2 * r_idx].imshow(img, cmap="gray") + axes[0, 2 * r_idx].set_title(f"{reconstruction} (u)") + axes[0, 2 * r_idx].axis("off") + else: + axes[0, 2 * r_idx].set_title(f"{reconstruction} (u)") + axes[0, 2 * r_idx].text( + 0.5, + 0.5, + "MISSING", + transform=axes[0, 2 * r_idx].transAxes, + fontsize=12, + color="red", + ha="center", + ) + axes[0, 2 * r_idx].axis("off") + + else: + axes[0, r_idx].axis("on") + axes[0, r_idx].xaxis.set_visible(False) + axes[0, r_idx].set_yticks([]) + axes[0, r_idx].set_ylabel("BaseDistortion", fontsize=12) + result_file_names_corr = glob.glob( os.path.join( result_folder, - f"image_{dataset_name}_{sample_name}_{distortion}*_{reconstruction}_corrected.tiff", + f"image_{dataset_name}_{sample_name}_BaseDistortion*_{reconstruction}_corrected.tiff", ) ) if len(result_file_names_corr) > 0: img = imread(result_file_names_corr[0]).squeeze() if len(img.shape) == 3: - print("Warning: image has 3 dimensions: [img.shape]") + print(f"Warning: image has 3 dimensions: {img.shape}") print(result_file_names_corr[0]) img = img[0, ...] - axes[d_idx + 1, 2 * r_idx + 1].imshow(img, cmap="gray") - axes[d_idx + 1, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") - axes[d_idx + 1, 2 * r_idx + 1].axis("off") + axes[0, 2 * r_idx + 1].imshow(img, cmap="gray") + axes[0, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") + axes[0, 2 * r_idx + 1].axis("off") else: - axes[d_idx + 1, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") - axes[d_idx + 1, 2 * r_idx + 1].text( + axes[0, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") + axes[0, 2 * r_idx + 1].text( 0.5, 0.5, "MISSING", - transform=axes[d_idx + 1, 2 * r_idx].transAxes, + transform=axes[0, 2 * r_idx + 1].transAxes, fontsize=12, color="red", ha="center", ) - axes[d_idx + 1, 2 * r_idx + 1].axis("off") - - result_file_names_uncorr = glob.glob( - os.path.join( - result_folder, - f"image_{dataset_name}_{sample_name}_{distortion}*_{reconstruction}_uncorrected.tiff", + axes[0, 2 * r_idx + 1].axis("off") + + # reconstruction methods in columns, distortions in rows + for d_idx, distortion in enumerate( + [dist for dist in distortion_names if dist != "BaseDistortion"] + ): + print("\t", distortion) + for r_idx, reconstruction in enumerate(reconstruction_names_part): + print("\t\t", reconstruction) + result_file_names_corr = glob.glob( + os.path.join( + result_folder, + f"image_{dataset_name}_{sample_name}_{distortion}*_{reconstruction}_corrected.tiff", + ) ) - ) - if len(result_file_names_uncorr) > 0: - img_u = imread(result_file_names_uncorr[0]).squeeze() if len(result_file_names_corr) > 0: - if (img == img_u).all(): - print( - f"Warning: corrected and uncorrected images are the same for {dataset_name}, {distortion}, {reconstruction}" - ) - if len(img_u.shape) == 3: - print(f"Warning: image has 3 dimensions: {img_u.shape}") - print(result_file_names_uncorr[0]) - img_u = img_u[0, ...] - axes[d_idx + 1, 2 * r_idx].imshow(img_u, cmap="gray") - axes[d_idx + 1, 2 * r_idx].set_title(f"{reconstruction} (u)") - axes[d_idx + 1, 2 * r_idx].axis("off") - else: - axes[d_idx + 1, 2 * r_idx].set_title(f"{reconstruction} (u)") - axes[d_idx + 1, 2 * r_idx].text( - 0.5, - 0.5, - "MISSING", - transform=axes[d_idx + 1, 2 * r_idx + 1].transAxes, - fontsize=12, - color="red", - ha="center", + img = imread(result_file_names_corr[0]).squeeze() + if len(img.shape) == 3: + print("Warning: image has 3 dimensions: [img.shape]") + print(result_file_names_corr[0]) + img = img[0, ...] + axes[d_idx + 1, 2 * r_idx + 1].imshow(img, cmap="gray") + axes[d_idx + 1, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") + axes[d_idx + 1, 2 * r_idx + 1].axis("off") + else: + axes[d_idx + 1, 2 * r_idx + 1].set_title(f"{reconstruction} (c)") + axes[d_idx + 1, 2 * r_idx + 1].text( + 0.5, + 0.5, + "MISSING", + transform=axes[d_idx + 1, 2 * r_idx].transAxes, + fontsize=12, + color="red", + ha="center", + ) + axes[d_idx + 1, 2 * r_idx + 1].axis("off") + + result_file_names_uncorr = glob.glob( + os.path.join( + result_folder, + f"image_{dataset_name}_{sample_name}_{distortion}*_{reconstruction}_uncorrected.tiff", + ) ) - axes[d_idx + 1, 2 * r_idx].axis("off") - - if r_idx == 0: - axes[d_idx + 1, 0].axis("on") - axes[d_idx + 1, 0].xaxis.set_visible(False) - axes[d_idx + 1, 0].set_yticks([]) - if len(distortion) > 30: - # find the latest capital letter between second and 30th character and split there - for c in distortion[1:30]: - if c.isupper(): - cap_letter_idx = distortion[1:30].find(c) - if cap_letter_idx != -1: - str_distortion = ( - distortion[: cap_letter_idx + 1] - + "\n" - + distortion[cap_letter_idx + 1 :] - ) - else: - str_distortion = distortion[:30] + "\n" + distortion[30:50] + if len(result_file_names_uncorr) > 0: + img_u = imread(result_file_names_uncorr[0]).squeeze() + if len(result_file_names_corr) > 0: + if (img == img_u).all(): + print( + f"Warning: corrected and uncorrected images are the same for {dataset_name}, {distortion}, {reconstruction}" + ) + if len(img_u.shape) == 3: + print(f"Warning: image has 3 dimensions: {img_u.shape}") + print(result_file_names_uncorr[0]) + img_u = img_u[0, ...] + axes[d_idx + 1, 2 * r_idx].imshow(img_u, cmap="gray") + axes[d_idx + 1, 2 * r_idx].set_title(f"{reconstruction} (u)") + axes[d_idx + 1, 2 * r_idx].axis("off") else: - str_distortion = distortion - axes[d_idx + 1, 0].set_ylabel(str_distortion, fontsize=12) - - plt.tight_layout() - plt.savefig( - os.path.join( - "/home/melanie.dohmen/ArtifactLab/reports", - f"summary_20260616_{dataset_name}_{sample_name}_part_{part + 1}.png", + axes[d_idx + 1, 2 * r_idx].set_title(f"{reconstruction} (u)") + axes[d_idx + 1, 2 * r_idx].text( + 0.5, + 0.5, + "MISSING", + transform=axes[d_idx + 1, 2 * r_idx + 1].transAxes, + fontsize=12, + color="red", + ha="center", + ) + axes[d_idx + 1, 2 * r_idx].axis("off") + + if r_idx == 0: + axes[d_idx + 1, 0].axis("on") + axes[d_idx + 1, 0].xaxis.set_visible(False) + axes[d_idx + 1, 0].set_yticks([]) + if len(distortion) > 30: + # find the latest capital letter between second and 30th character and split there + for c in distortion[1:30]: + if c.isupper(): + cap_letter_idx = distortion[1:30].find(c) + if cap_letter_idx != -1: + str_distortion = ( + distortion[: cap_letter_idx + 1] + + "\n" + + distortion[cap_letter_idx + 1 :] + ) + else: + str_distortion = distortion[:30] + "\n" + distortion[30:50] + else: + str_distortion = distortion + axes[d_idx + 1, 0].set_ylabel(str_distortion, fontsize=12) + + plt.tight_layout() + plt.savefig( + os.path.join( + result_folder, + f"summary_20260616_{dataset_name}_{sample_name}_part_{part + 1}.png", + ) ) - ) ###### # Create Plots for each distortion and sample ###### +def find_all(string, substring): + indices = [] + start = 0 + while True: + start = string.find(substring, start) + if start == -1: + break + indices.append(start) + start += len(substring) # Move past the last found substring + return indices + -for dataset_name, sample_names in samples.items(): - for sample_name in sample_names: - print("looping over ", dataset_name, " ", sample_name) +create_distortion_summary = True - for d_idx, distortion in enumerate( - [dist for dist in distortion_names if dist != "BaseDistortion"] - ): - print("\t", distortion) +if create_distortion_summary: + for dataset_name, sample_names in samples.items(): + for sample_name in sample_names: + print("looping over ", dataset_name, " ", sample_name) - results_for_distortion = glob.glob( - os.path.join( - result_folder, f"image_{dataset_name}_{sample_name}_{distortion}*.tiff" + for d_idx, distortion in enumerate( + [dist for dist in distortion_names if dist != "BaseDistortion"] + ): + results_for_distortion = glob.glob( + os.path.join( + result_folder, f"image_{dataset_name}_{sample_name}_{distortion}*.tiff" + ) ) - ) - nr_rows = np.ceil(np.sqrt(len(results_for_distortion))) - nr_cols = np.ceil(len(results_for_distortion) / nr_rows) - fig, axes = plt.subplots( - int(nr_rows), int(nr_cols), figsize=(3 * nr_cols, 3 * nr_rows), squeeze=False - ) - fig.suptitle(f"{dataset_name} {sample_name}") + # find reference: + reference = glob.glob( + os.path.join( + result_folder, f"image_{dataset_name}_{sample_name}_reference.tiff" + ) + ) - # sort results: - results_for_distortion_sorted = sorted( - results_for_distortion, key=lambda x: (x.split("_")[-2], x.split("_")[-1]) - ) + nr_rows = int(np.ceil(np.sqrt(len(results_for_distortion)))) + len(reference) + if nr_rows > 1: + print( + f"For {distortion} found {len(results_for_distortion)} results and {len(reference)} reference images" + ) + nr_cols = int(np.ceil((len(results_for_distortion) + len(reference)) / nr_rows)) - # set BaseDistortion and CG first: - results_for_distortion_sorted = sorted( - results_for_distortion_sorted, - key=lambda x: ( - x.split("_")[-2] != "BaseDistortion", - x.split("_")[-2] != "conjugate-gradient", - ), - ) + fig, axes = plt.subplots( + int(nr_rows), + int(nr_cols), + figsize=(3 * nr_cols, 3 * nr_rows), + squeeze=False, + ) - for r_idx, result_file_name in enumerate(results_for_distortion_sorted): - # split filename to get reconstruction name and corrected/uncorrected - reconstruction = result_file_name.split("_")[-2] - corrected = result_file_name.split("_")[-1].split(".")[0] - - img = imread(result_file_name).squeeze() - if len(img.shape) == 3: - print(f"Warning: image has 3 dimensions: {img.shape}") - print(result_file_name) - img = img[0, ...] - axes[r_idx // nr_cols, r_idx % nr_cols].imshow(img, cmap="gray") - axes[r_idx // nr_cols, r_idx % nr_cols].set_title(f"{reconstruction} (c)") - axes[r_idx // nr_cols, r_idx % nr_cols].axis("off") - - plt.tight_layout() - plt.savefig( - os.path.join( - "/home/melanie.dohmen/ArtifactLab/reports", - f"summary_20260616_{dataset_name}_{sample_name}_{distortion}.png", - ) - ) + # sort results: + results_for_distortion_sorted = reference + sorted( + results_for_distortion, key=lambda x: (x.split("_")[-2], x.split("_")[-1]) + ) + + for r_idx, result_file_name in enumerate(results_for_distortion_sorted): + # split filename to get reconstruction name and corrected/uncorrected and parameters + # to add details to each result + reconstruction = result_file_name.split("_")[-2] + corrected = result_file_name.split("_")[-1].split(".")[0] + parameters_str = result_file_name.split("_")[-3][len(distortion) :] + " " + # example parameters: 'k=0.15w=2b=35 ' + parameters_indices = find_all(parameters_str, "=") + [len(parameters_str)] + # example parameters_indices = [1, 7, 10] + [14] + parameters_list = [ + parameters_str[i - 1 : parameters_indices[i_idx + 1] - 1] + for i_idx, i in enumerate(parameters_indices[:-1]) + ] + # example parameters_list = ['k=0.15', 'w=2', 'b=35'] + parameters = "\n" + " ".join(parameters_list) + img = imread(result_file_name).squeeze() + if len(img.shape) == 3: + print(f"Warning: image has 3 dimensions: {img.shape}") + print(result_file_name) + img = img[0, ...] + axes[r_idx // nr_cols, r_idx % nr_cols].imshow(img, cmap="gray") + if r_idx < len(reference): + axes[r_idx // nr_cols, r_idx % nr_cols].set_title( + f"{distortion}\n{dataset_name}{sample_name}\n(reference)" + ) + else: + axes[r_idx // nr_cols, r_idx % nr_cols].set_title( + f"{reconstruction} ({corrected[0]}){parameters}" + ) + axes[r_idx // nr_cols, r_idx % nr_cols].axis("off") + + # remove axis for empty subplots + for r_idx in range(len(results_for_distortion_sorted), nr_rows * nr_cols): + axes[r_idx // nr_cols, r_idx % nr_cols].axis("off") + + plt.tight_layout() + plt.savefig( + os.path.join( + result_folder, + f"summary_20260616_{dataset_name}_{sample_name}_{distortion}.png", + ) + ) + + else: + print("No results for ", distortion) From 43bc1b2df05df770814fa32cd1511d3368ed6f96 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Wed, 8 Jul 2026 22:31:16 +0000 Subject: [PATCH 20/22] test multiple distortion with different result paths --- examples/test_dist_params.py | 42 +++++++++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/examples/test_dist_params.py b/examples/test_dist_params.py index 5f3a354..039d038 100644 --- a/examples/test_dist_params.py +++ b/examples/test_dist_params.py @@ -14,7 +14,7 @@ "data": { "fastmri_knee": "/home/melanie.dohmen/ArtifactLab/data/singlecoil_val", "oasis": "/home/melanie.dohmen/ArtifactLab/data/oasis", - "fastmri_brain": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_multicoil_brain_test", + "fastmri_brain": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_multicoil_brain_train0", "cmrxrecon": "/home/melanie.dohmen/ArtifactLab/data/CMRxRecon", "fastmri_prostate": "/home/melanie.dohmen/ArtifactLab/data/fastMRI_prostate_T2_IDS_001_020", }, @@ -28,6 +28,7 @@ "conjugate-gradient", ], "add_N4Correction": False, + "resolution_reduction_factors": [], "num_samples": 1, "verbose": True, "results_dir": results_dir, @@ -53,16 +54,51 @@ # "center_y_fraction": [-0.1, -0.1, -0.1, -0.1, -0.1,-0.1,-0.1,-0.1, -0.1,-0.1,-0.1,-0.1,], # "edge_gain": [0.05,0.05,0.05,0.05, 0.1, 0.1, 0.1, 0.1, 0.5, 0.5, 0.5, 0.5, ], # }}, + # { + # "CartesianUndersamplingEquispacedZeroACS": { + # "keep_fraction": [0.1, 0.25, 0.5, 0.75, 0.85, 0.9, 0.95, 0.98] + # }, + # }, { - "CartesianUndersamplingEquispacedZeroACS": { - "keep_fraction": [0.1, 0.25, 0.5, 0.75, 0.85, 0.9, 0.95, 0.98] + "AnisotropicLP": { + "kx_radius_fraction": [0.1, 0.25, 0.5, 0.75, 0.85, 0.9, 0.95, 1.0], + "ky_radius_fraction": [1.0, 0.95, 0.6, 0.85, 0.75, 0.5, 0.25, 0.1], }, }, + { + "HannTaperLP": { + "radius_fraction": [0.1, 0.1, 0.5, 0.5, 0.9, 0.9, 1.0, 1.0], # 0.35, + "transition_fraction": [0.2, 0.6, 0.2, 0.6, 0.2, 0.6, 0.2, 0.6], # 0.4, + } + }, + { + "KaiserTaperLP": { + "radius_fraction": [0.1, 0.1, 0.5, 0.5, 0.9, 0.9, 1.0, 1.0], # 0.35, + "transition_fraction": [0.2, 0.6, 0.2, 0.6, 0.2, 0.6, 0.2, 0.6], # 0.4, + "beta": [8.6, 8.6, 8.6, 8.6, 2.0, 2.0, 8.6, 8.6], # 8.6 + } + }, + { + "GaussianNoise": { + "sigma": [0.00001, 0.0001, 0.001, 0.01, 0.1, 1.0], # [0.00001] + } + }, + { + "IsotropicLP": { + "radius_fraction": [0.1, 0.25, 0.5, 0.75, 0.85, 0.9, 0.95, 1.0], # 0.1 + } + }, + { + "RadialHighPassEmphasis": { + "alpha": [0.1, 0.2, 0.4, 0.5, 0.75, 0.9, 1.0], # 0.4 + } + }, ] for d_idx, distortion_dict in enumerate(distortions): for distortion_name, dist_params in distortion_dict.items(): nr_param_values = len(dist_params[list(dist_params.keys())[0]]) + config["distortions"] = [] for v_idx in range(nr_param_values): single_value_distortion_dict = { distortion_name: { From 7413f7d421d1c6d41f6accd281db73fad7a9217c Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Wed, 8 Jul 2026 22:31:45 +0000 Subject: [PATCH 21/22] sort order of samples for reproducibility --- mri_recon/utils/prostate_adaptor.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mri_recon/utils/prostate_adaptor.py b/mri_recon/utils/prostate_adaptor.py index a21184e..4c7a8a3 100644 --- a/mri_recon/utils/prostate_adaptor.py +++ b/mri_recon/utils/prostate_adaptor.py @@ -22,7 +22,9 @@ def __init__( def get_image_data(self) -> tuple[np.ndarray, list[str]]: image_result_list = [] sample_name_list = [] - for sample_idx, filename in enumerate(glob.glob(os.path.join(self.data_path, "*.h5"))): + for sample_idx, filename in enumerate( + sorted(glob.glob(os.path.join(self.data_path, "*.h5"))) + ): if (self.num_samples is not None) and (sample_idx >= self.num_samples): break try: From aefad4f55d678c33964ff9ea7cb568b2b3589fc6 Mon Sep 17 00:00:00 2001 From: Melanie Dohmen Date: Wed, 8 Jul 2026 22:54:06 +0000 Subject: [PATCH 22/22] optionally do not overwrite existing results --- examples/config.yaml | 1 + examples/run_all.py | 114 ++++++++++++++++++++++++++----------------- 2 files changed, 70 insertions(+), 45 deletions(-) diff --git a/examples/config.yaml b/examples/config.yaml index 4d2d3ae..b360198 100644 --- a/examples/config.yaml +++ b/examples/config.yaml @@ -96,4 +96,5 @@ reconstruction_algorithms: add_N4Correction: false num_samples: 1 verbose: true +overwrite: false results_dir: "/home/melanie.dohmen/ArtifactLab/reports/experiments_run1" diff --git a/examples/run_all.py b/examples/run_all.py index 2a68e5f..5ccf6ce 100644 --- a/examples/run_all.py +++ b/examples/run_all.py @@ -111,13 +111,14 @@ def get_measurement_sample( and "coil_maps" in sample_batch[2] else None ) - # centered k-space data, shape: (B, 2, n_timepoints, (n_coils), H, W) dtype: float32 - y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) + # select the center timepoint in order to simplify the evaluation of the reconstruction algorithms center_time_point = y.shape[2] // 2 y = y[:, :, center_time_point, ...] - y_centered = y_centered[:, :, center_time_point, ...] + + # centered k-space data, shape: (B, 2, H, W) dtype: float32 + y_centered = fastmri_measurement_to_oasis_kspace(y, coil_maps=coil_maps, device=run_device) x = x[:, :, center_time_point, ...] elif dataset_name == "fastmri_prostate": @@ -246,6 +247,22 @@ def run_all(config) -> None: y_distorted = distortion.A(y) + # save reference and distorted k-space for debugging purposes + imwrite( + os.path.join( + config["results_dir"], + f"kspace_{dataset_name}_{sample_name}_reference.tiff", + ), + _kspace_to_log_magnitude(y).numpy(), + ) + imwrite( + os.path.join( + config["results_dir"], + f"kspace_{dataset_name}_{sample_name}_{distortion_name_with_params}.tiff", + ), + _kspace_to_log_magnitude(y_distorted).numpy(), + ) + physics_distorted = DistortedKspaceMultiCoilMRI( distortion, img_size=y.shape[-2:], @@ -254,8 +271,21 @@ def run_all(config) -> None: ) for reconstructor_name in config["reconstruction_algorithms"]: + + corrected_reconstructed_image_filename = os.path.join( + config["results_dir"], + f"image_{dataset_name}_{sample_name}_{distortion_name_with_params}_{reconstructor_name}_corrected.tiff", + ) + uncorrected_reconstructed_image_filename = os.path.join( + config["results_dir"], + f"image_{dataset_name}_{sample_name}_{distortion_name_with_params}_{reconstructor_name}_uncorrected.tiff", + ) + # only run on reconstructors, that use the fastmri-like k-space - if not uses_oasis_centered_path(reconstructor_name): + if (not uses_oasis_centered_path(reconstructor_name)) and (config["overwrite"] or + not os.path.exists(corrected_reconstructed_image_filename)) and ( + not os.path.exists(uncorrected_reconstructed_image_filename)): + print(f"\t\t{reconstructor_name} ...") start = datetime.now() if compatible_dataset_with_reconstructor( @@ -268,22 +298,7 @@ def run_all(config) -> None: verbose=config["verbose"], ).to(device) - # save reference and distorted k-space for debugging purposes - imwrite( - os.path.join( - config["results_dir"], - f"kspace_{dataset_name}_{sample_name}_reference.tiff", - ), - _kspace_to_log_magnitude(y).numpy(), - ) - imwrite( - os.path.join( - config["results_dir"], - f"kspace_{dataset_name}_{sample_name}_{distortion_name_with_params}.tiff", - ), - _kspace_to_log_magnitude(y_distorted).numpy(), - ) - + # actual reconstruction with the selected reconstructor try: x_uncorrected = reconstructor(y_distorted, physics_clean) @@ -302,17 +317,11 @@ def run_all(config) -> None: # save reconstructed images imwrite( - os.path.join( - config["results_dir"], - f"image_{dataset_name}_{sample_name}_{distortion_name_with_params}_{reconstructor_name}_uncorrected.tiff", - ), + uncorrected_reconstructed_image_filename, convert_image_for_save(x_uncorrected), ) imwrite( - os.path.join( - config["results_dir"], - f"image_{dataset_name}_{sample_name}_{distortion_name_with_params}_{reconstructor_name}_corrected.tiff", - ), + corrected_reconstructed_image_filename, convert_image_for_save(x_corrected), ) print(f"\t\t... done in {datetime.now() - start}") @@ -343,11 +352,41 @@ def run_all(config) -> None: distortion.A(torch.fft.fftshift(y_centered, dim=(-1, -2))), dim=(-2, -1) ) + # save reference and distorted k-space for debugging purposes + imwrite( + os.path.join( + config["results_dir"], + f"kspace_centered_{dataset_name}_{sample_name}_reference.tiff", + ), + _kspace_to_log_magnitude(y_centered).numpy(), + ) + imwrite( + os.path.join( + config["results_dir"], + f"kspace_centered_{dataset_name}_{sample_name}_{distortion_name_with_params}.tiff", + ), + _kspace_to_log_magnitude(y_distorted).numpy(), + ) + + physics_distorted = OasisCenteredFFTPhysics(distortion) for reconstructor_name in config["reconstruction_algorithms"]: + + corrected_reconstructed_image_filename = os.path.join( + config["results_dir"], + f"image_{dataset_name}_{sample_name}_{distortion_name_with_params}_{reconstructor_name}_corrected.tiff", + ) + uncorrected_reconstructed_image_filename = os.path.join( + config["results_dir"], + f"image_{dataset_name}_{sample_name}_{distortion_name_with_params}_{reconstructor_name}_uncorrected.tiff", + ) # skip all reconstructors, that don't use the oasis-centered path - if uses_oasis_centered_path(reconstructor_name): + if (uses_oasis_centered_path(reconstructor_name) and + (config["overwrite"] or + (not os.path.exists(corrected_reconstructed_image_filename) and + not os.path.exists(uncorrected_reconstructed_image_filename)))): + print(f"\t\t{reconstructor_name} ...") start = datetime.now() if compatible_dataset_with_reconstructor( @@ -360,22 +399,7 @@ def run_all(config) -> None: verbose=config["verbose"], ).to(device) - # save reference and distorted k-space for debugging purposes - imwrite( - os.path.join( - config["results_dir"], - f"kspace_centered_{dataset_name}_{sample_name}_reference.tiff", - ), - _kspace_to_log_magnitude(y_centered).numpy(), - ) - imwrite( - os.path.join( - config["results_dir"], - f"kspace_centered_{dataset_name}_{sample_name}_{distortion_name_with_params}.tiff", - ), - _kspace_to_log_magnitude(y_distorted).numpy(), - ) - + # actual reconstruction with the algo being evaluated try: x_uncorrected = reconstructor(y_distorted, physics_clean)