diff --git a/src/mri/cli/reconstruct.py b/src/mri/cli/reconstruct.py index 2c9df3e4..95bf0537 100644 --- a/src/mri/cli/reconstruct.py +++ b/src/mri/cli/reconstruct.py @@ -3,9 +3,10 @@ from mri.io.output import save_data from mri.cli.utils import raw_config, traj_config, setup_hydra_config, get_outdir_path from mri.operators.fourier.utils import discard_frequency_outliers -from mrinufft.io.utils import add_phase_to_kspace_with_shifts +from mrinufft.io.utils import add_phase_to_kspace_with_shifts, remove_extra_kspace_samples from pymrt.recipes.coils import compress_svd from mri.reconstructors import SelfCalibrationReconstructor +from mri.reconstructors.ggrappa import do_grappa_and_append_data, GRAPPA_RECON_AVAILABLE import json import numpy as np @@ -20,7 +21,8 @@ def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int, debug: int, - obs_reader, traj_reader, fourier, output_filename: str = "dc_adjoint.nii"): + obs_reader, traj_reader, fourier, grappa_recon=None, output_filename: str = "dc_adjoint.nii", + return_data=False): """ Reconstructs an image using the adjoint operator. @@ -48,6 +50,8 @@ def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int, It can be: 1) *.pkl / *.mat: Holds the reconstructed results saved in dictionary as `recon`. 2) *.nii : NIFTI file holding the reconstructed images. + grappa_af: Union[list[int], tuple[int, ...]], optional default 1 + The acceleration factor for the GRAPPA reconstruction. Returns ------- @@ -86,7 +90,6 @@ def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int, for size in traj_params['img_size'] ]) log.info(f"Trajectory Parameters: {traj_params}") - kspace_loc = shots.reshape(-1, traj_params["dimension"]) data_header["shifts"] = data_header['shifts'][:traj_params["dimension"]] normalized_shifts = ( np.array(data_header["shifts"]) @@ -94,25 +97,48 @@ def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int, * np.array(traj_params["img_size"]) / 1000 ) - if kspace_loc.max() > 0.5 or kspace_loc.min() < 0.5: - log.warn(f"K-space locations are above the unity range, discarding the outlier data") - if data_header["type"] == "retro_recon": - kspace_loc = discard_frequency_outliers(kspace_loc) - kspace_data = np.squeeze(raw_data) - else: - kspace_loc, kspace_data = discard_frequency_outliers(kspace_loc, np.squeeze(raw_data)) - kspace_data = kspace_data.astype(np.complex64) - kspace_loc = kspace_loc.astype(np.float32) + kspace_data = np.squeeze(raw_data).astype(np.complex64) + kspace_loc = shots.reshape(-1, traj_params["dimension"]).astype(np.float32) + kspace_data = remove_extra_kspace_samples(kspace_data, shots.shape[1]) + kspace_data = kspace_data.reshape(kspace_data.shape[0], -1) log.info(f"Phase shifting raw data for Normalized shifts: {normalized_shifts}") kspace_data = add_phase_to_kspace_with_shifts( - kspace_data, kspace_loc, normalized_shifts + kspace_data, kspace_loc.reshape(-1, traj_params["dimension"]), normalized_shifts ) + try: + af_string = data_header['trajectory_name'].split('_G')[1].split('_')[0].split('x') + if len(af_string) > 1 and 'd' in af_string[1]: + af_caipi = af_string[1].split('d') + af_string[1] = af_caipi[0] + grappa_recon.keywords['delta'] = int(af_caipi[1]) + grappa_recon.keywords['af'] = tuple([int(float(af)) for af in af_string]) + except: + grappa_recon.keywords['af'] = (1, ) + grappa_recon.keywords['delta'] = 0 + if grappa_recon is not None and np.prod(grappa_recon.keywords['af'])>1: + log.info("Performing GRAPPA Reconstruction: AF: %s", af_string) + log.info("GRAPPA AF: %s", grappa_recon.keywords['af']) + kspace_loc, kspace_data = do_grappa_and_append_data( + kspace_loc, + kspace_data, + traj_params, + grappa_recon, + acs=data_header["acs"], # Pass ACS if read in data (external) + ) if coil_compress != -1: + log.info("Compressing coils") kspace_data = np.ascontiguousarray(compress_svd( kspace_data, k_svd=coil_compress, coil_axis=0 - )) + )).astype(np.complex64) + if kspace_loc.max() > 0.5 or kspace_loc.min() < 0.5: + log.warn(f"K-space locations are above the unity range, discarding the outlier data") + if data_header["type"] == "retro_recon": + kspace_loc = discard_frequency_outliers(kspace_loc) + kspace_data = np.squeeze(raw_data) + else: + kspace_loc, kspace_data = discard_frequency_outliers(kspace_loc, kspace_data) fourier.keywords['smaps'] = partial( fourier.keywords['smaps'], kspace_data=kspace_data, @@ -127,6 +153,7 @@ def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int, 'density_comp': fourier_op.impl.density, 'traj_params': traj_params, 'data_header': data_header, + 'kspace_loc': kspace_loc, } save_data_hydra('smaps.nii', fourier_op.impl.smaps) if coil_compress != -1: @@ -140,13 +167,15 @@ def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int, log.info("Saving DC Adjoint") data_header['traj_params'] = traj_params save_data_hydra(output_filename, dc_adjoint, data_header) - return dc_adjoint, (fourier_op, kspace_data, traj_params, data_header) + if return_data: + return dc_adjoint, (fourier_op, kspace_data, traj_params, data_header) def recon(obs_file: str, traj_file: str, mu: float, num_iterations: int, coil_compress: str|int, algorithm: str, debug: int, obs_reader, traj_reader, fourier, linear, sparsity, - output_filename: str = "recon.nii", remove_dc_for_recon: bool = True, validation_recon: np.ndarray = None, metrics: dict = None): + output_filename: str = "recon.nii", remove_dc_for_recon: bool = True, validation_recon: np.ndarray = None, metrics: dict = None, + grappa_recon=None): """Reconstructs an MRI image using the given parameters. Parameters @@ -193,7 +222,9 @@ def recon(obs_file: str, traj_file: str, mu: float, num_iterations: int, coil_co obs_reader, traj_reader, fourier, + grappa_recon=grappa_recon, output_filename='dc_adj_' + output_filename, + return_data=True, ) fourier_op, kspace_data, traj_params, data_header = additional_data if remove_dc_for_recon: @@ -246,6 +277,7 @@ def recon(obs_file: str, traj_file: str, mu: float, num_iterations: int, coil_co "_self_", {"fourier": "gpu"}, {"fourier/density_comp": "pipe"}, + {"grappa_recon": "disable"} if GRAPPA_RECON_AVAILABLE else {}, {"fourier/smaps": "low_frequency"}, ], name="dc_adjoint", @@ -257,11 +289,13 @@ def recon(obs_file: str, traj_file: str, mu: float, num_iterations: int, coil_co algorithm="pogm", num_iterations=30, coil_compress=10, + mu=1e-7, debug=1, hydra_defaults=[ "_self_", {"fourier": "gpu"}, {"fourier/density_comp": "pipe"}, + {"grappa_recon": "disable"} if GRAPPA_RECON_AVAILABLE else {}, {"fourier/smaps": "low_frequency"}, {"linear": "gpu"}, {"sparsity": "weighted_sparse"}, @@ -276,10 +310,12 @@ def recon(obs_file: str, traj_file: str, mu: float, num_iterations: int, coil_co algorithm="pogm", num_iterations=30, coil_compress=5, + mu=1e-7, debug=1, hydra_defaults=[ "_self_", {"fourier": "gpu_lowmem"}, + {"grappa_recon": "disable"} if GRAPPA_RECON_AVAILABLE else {}, {"fourier/density_comp": "pipe_lowmem"}, {"fourier/smaps": "low_frequency"}, ], diff --git a/src/mri/cli/utils.py b/src/mri/cli/utils.py index 7092c7e7..e51c69d2 100644 --- a/src/mri/cli/utils.py +++ b/src/mri/cli/utils.py @@ -12,6 +12,18 @@ from modopt.opt.linear import Identity import os +try: + from ggrappa.grappaND import GRAPPA_Recon + grappa_config = builds( + GRAPPA_Recon, + zen_exclude=["sig", "acs", "isGolfSparks", "quiet", "af"], + zen_partial=True, + populate_full_signature=True, + ) + grappa_store = store(group="grappa_recon") + grappa_store(grappa_config, name="disable") +except: + pass raw_config = builds(read_arbgrad_rawdat, populate_full_signature=True, zen_partial=True) @@ -31,6 +43,7 @@ smaps_config = builds( get_smaps("low_frequency"), populate_full_signature=True, + blurr_factor=30.0, # We estimate density, with separate args. It is passed by compute_smaps in mri-nufft zen_exclude=["density"], zen_partial=True, diff --git a/src/mri/reconstructors/ggrappa.py b/src/mri/reconstructors/ggrappa.py new file mode 100644 index 00000000..c0c030cb --- /dev/null +++ b/src/mri/reconstructors/ggrappa.py @@ -0,0 +1,27 @@ +import numpy as np + +GRAPPA_RECON_AVAILABLE = False +try: + from ggrappa.grappaND import GRAPPA_Recon + import torch + from ggrappa.utils import get_cart_portion_sparkling, get_grappa_filled_data_and_loc + GRAPPA_RECON_AVAILABLE = True +except: + pass + + +def do_grappa_and_append_data(kspace_loc, kspace_data, traj_params, grappa_maker, acs=None): + kspace_shots = kspace_loc.reshape(traj_params['num_shots'], -1, traj_params['dimension']) + if not GRAPPA_RECON_AVAILABLE: + raise ValueError("GRAPPA is not available") + gridded_center = get_cart_portion_sparkling(kspace_shots, traj_params, kspace_data) + grappa_recon, grappa_kernel = grappa_maker( + sig=torch.tensor(gridded_center).permute(0, 2, 3, 1), + acs=torch.tensor(acs).permute(0, 2, 3, 1) if acs is not None else None, + isGolfSparks=True, + ) + grappa_recon = grappa_recon.permute(0, 3, 1, 2).numpy() + extra_loc, extra_data = get_grappa_filled_data_and_loc(gridded_center, grappa_recon, traj_params) + kspace_loc = np.concatenate([kspace_loc, extra_loc], axis=0) + kspace_data = np.hstack([kspace_data, extra_data]) + return kspace_loc, kspace_data \ No newline at end of file