diff --git a/warpkit/distortion.py b/warpkit/distortion.py index c685478..34ee5d0 100644 --- a/warpkit/distortion.py +++ b/warpkit/distortion.py @@ -3,8 +3,9 @@ import nibabel as nib import numpy as np -from warpkit.unwrap import unwrap_and_compute_field_maps +from warpkit.unwrap import unwrap_phase_data, compute_field_maps from warpkit.utilities import ( + check_affines, displacement_maps_to_field_maps, field_maps_to_displacement_maps, invert_displacement_maps, @@ -69,37 +70,32 @@ def medic( nib.Nifti1Image Field maps in Hz (undistorted space) """ - # make sure affines/shapes are all correct - for p1, m1 in zip(phase, mag): - for p2, m2 in zip(phase, mag): - if not ( - np.allclose(p1.affine, p2.affine, rtol=1e-3, atol=1e-3) - and np.allclose(m1.affine, m2.affine, rtol=1e-3, atol=1e-3) - and p1.shape == p2.shape - and m1.shape == m2.shape - ): - print(p1.affine, p2.affine) - print(p1.affine - p2.affine) - print(p1.shape, p2.shape) - print(m1.affine, m2.affine) - print(m1.affine - m2.affine) - print(m1.shape, m2.shape) - raise ValueError("Affines and shapes must match") + check_affines(phase, mag) # unwrap phase and compute field maps try: - field_maps_native = unwrap_and_compute_field_maps( + unwrapped, new_masks = unwrap_phase_data( phase, mag, TEs, + automask=True, border_size=border_size, - border_filt=border_filt, - svd_filt=svd_filt, frames=frames, n_cpus=n_cpus, debug=debug, wrap_limit=wrap_limit, ) + field_maps_native = compute_field_maps( + unwrapped=unwrapped, + new_masks=new_masks, + img=phase[0], + mag=mag, + TEs=TEs, + border_filt=border_filt, + svd_filt=svd_filt, + frames=frames, + n_cpus=n_cpus, + ) except IndexError as e: raise IndexError( "An IndexError was encountered while unwrapping phase images. " diff --git a/warpkit/scripts/warpkit_unwrap.py b/warpkit/scripts/warpkit_unwrap.py new file mode 100644 index 0000000..3e3c43f --- /dev/null +++ b/warpkit/scripts/warpkit_unwrap.py @@ -0,0 +1,199 @@ +"""A workflow to unwrap multi-echo phase data.""" + +import argparse +import json +import logging +from functools import partial +from pathlib import Path + +import nibabel as nib +import numpy as np + +from warpkit.distortion import unwrap_phase_data +from warpkit.scripts import epilog +from warpkit.utilities import setup_logging + + +def main(): + """Build parser object and run workflow.""" + + def _path_exists(path, parser): + """Ensure a given path exists.""" + if path is None or not Path(path).exists(): + raise parser.error(f"Path does not exist: <{path}>.") + return Path(path).absolute() + + def _is_file(path, parser): + """Ensure a given path exists and it is a file.""" + path = _path_exists(path, parser) + if not path.is_file(): + raise parser.error(f"Path should point to a file (or symlink of file): <{path}>.") + return path + + parser = argparse.ArgumentParser( + description="Unwrap multi-echo phase data", + epilog=f"{epilog} 12/09/2022", + ) + + IsFile = partial(_is_file, parser=parser) + + parser.add_argument( + "--magnitude", + nargs="+", + required=True, + metavar="FILE", + type=IsFile, + help="Magnitude data", + ) + parser.add_argument( + "--phase", + nargs="+", + required=True, + metavar="FILE", + type=IsFile, + help="Phase data", + ) + parser.add_argument( + "--metadata", + nargs="+", + required=True, + metavar="FILE", + type=IsFile, + help=( + "JSON sidecar for each echo. " + "Three fields are required: EchoTime, TotalReadoutTime, and PhaseEncodingDirection." + ), + ) + parser.add_argument( + "--out_prefix", + help="Prefix to output field maps and displacment maps.", + ) + parser.add_argument( + "-f", + "--noiseframes", + type=int, + default=0, + help=( + "Number of noise frames at the end of the run. " + "Noise frames will be removed before unwrapping is performed." + ), + ) + parser.add_argument( + "-n", + "--n_cpus", + type=int, + default=4, + help="Number of CPUs to use.", + ) + parser.add_argument( + "--debug", + action="store_true", + help="Debug mode", + ) + parser.add_argument( + "--wrap_limit", + action="store_true", + default=False, + help="Turns off some heuristics for phase unwrapping", + ) + + # parse arguments + args = parser.parse_args() + + # setup logging + setup_logging() + + # log arguments + logging.info(f"unwrap_phases: {args}") + kwargs = vars(args) + unwrap_phases(**kwargs) + + +def unwrap_phases( + *, + magnitude, + phase, + metadata, + out_prefix, + noiseframes, + border_size=5, + n_cpus=1, + debug=False, + wrap_limit=False, +): + """Unwrap multi-echo phase data. + + Parameters + ---------- + magnitude : list of str + List of magnitude data files. + phase : list of str + List of phase data files. + metadata : list of str + List of JSON sidecar files for each echo. + out_prefix : str + Prefix to output field maps and displacment maps. + noiseframes : int + Number of noise frames at the end of the run. + Noise frames will be removed before unwrapping is performed. + border_size : int, optional + Size of border in automask, by default 5 + n_cpus : int + Number of CPUs to use. + debug : bool + Debug mode. + wrap_limit : bool + Turns off some heuristics for phase unwrapping. + """ + # load magnitude and phase data + magnitude_imgs = [nib.load(m) for m in magnitude] + phase_imgs = [nib.load(p) for p in phase] + + # if noiseframes specified, remove them + if noiseframes > 0: + logging.info(f"Removing {noiseframes} noise frames from the end of the run...") + magnitude_imgs = [m.slicer[..., : -noiseframes] for m in magnitude_imgs] + phase_imgs = [p.slicer[..., : -noiseframes] for p in phase_imgs] + + # check if data is 4D or 3D + if phase_imgs[0].ndim == 3: + # convert data to 4D + phase_imgs = [ + nib.Nifti1Image(p.get_fdata()[..., np.newaxis], p.affine, p.header) for p in phase_imgs + ] + magnitude_imgs = [ + nib.Nifti1Image(m.get_fdata()[..., np.newaxis], m.affine, m.header) for m in magnitude_imgs + ] + else: + raise ValueError("Data must be 3D or 4D.") + + # get metadata + echo_times = [] + for json_file in metadata: + with open(json_file, "r") as fobj: + metadata_dict = json.load(fobj) + echo_times.append(metadata_dict["EchoTime"] * 1000) # convert TE from s to ms + + # Sort the echo times and data by echo time + echo_times, magnitude_imgs, phase_imgs = zip(*sorted(zip(echo_times, magnitude_imgs, phase_imgs))) + + # now run MEDIC's phase-unwrapping method + unwrapped_phases, _ = unwrap_phase_data( + phase=phase_imgs, + mag=magnitude_imgs, + TEs=echo_times, + automask=True, + border_size=border_size, + n_cpus=n_cpus, + debug=debug, + wrap_limit=wrap_limit, + ) + unwrapped_phases = [ + nib.Nifti1Image(ph, phase_imgs[0].affine, phase_imgs[0].header) for ph in unwrapped_phases + ] + + # save the fmaps and dmaps to file + logging.info("Saving field maps and displacement maps to file...") + for i_echo, unwrapped_phase in enumerate(unwrapped_phases): + unwrapped_phase.to_filename(f"{out_prefix}_echo-{i_echo + 1}_phase.nii.gz") + logging.info("Done.") diff --git a/warpkit/unwrap.py b/warpkit/unwrap.py index 5e9794b..47745dd 100644 --- a/warpkit/unwrap.py +++ b/warpkit/unwrap.py @@ -19,6 +19,7 @@ from .julia import JuliaContext from .model import weighted_regression from .utilities import ( + check_affines, corr2_coeff, create_brain_mask, get_largest_connected_component, @@ -170,22 +171,13 @@ def mcpc_3d_s( new_proposed_fieldmap, new_proposed_unwrapped_phases = get_dual_echo_fieldmap( new_proposed_phases, TEs, mags, mask ) - new_voxel_prop = ( - np.count_nonzero(new_proposed_fieldmap[voxel_mask] > 0) / new_proposed_fieldmap[voxel_mask].shape[0] - ) + # fit linear model to the proposed phases new_phase_fits = np.concatenate( (np.zeros((*new_proposed_unwrapped_phases.shape[:-1], 1)), new_proposed_unwrapped_phases), axis=-1 ) _, residuals_2, _, _, _ = np.polyfit(all_TEs, new_phase_fits[voxel_mask, :].T, 1, full=True) - # print(f"mean_proposed_fieldmap 1: {proposed_fieldmap[voxel_mask].mean()}") - # print(f"voxel_prop 1 : {voxel_prop}") - # print(f"mean_residuals 1: {residuals_1.mean()}") - # print(f"mean_phase_offset 1: {mean_phase_offset}") - # print(f"proposed_fieldmap 2: {new_proposed_fieldmap[voxel_mask].mean()}") - # print(f"voxel_prop 2: {new_voxel_prop}") - # print(f"mean_residuals 2: {residuals_2.mean()}") - # print(f"mean_phase_offset 2: {new_proposed_offset.mean()}") + if ( np.isclose(residuals_1.mean(), residuals_2.mean(), atol=1e-3, rtol=1e-3) and new_proposed_fieldmap[voxel_mask].mean() > 0 @@ -589,105 +581,43 @@ def svd_filtering( field_maps[new_masks[..., i_vol] > 0, i_vol] = recon_img[new_masks[..., i_vol] > 0, i_vol] -def unwrap_and_compute_field_maps( +def unwrap_phase_data( phase: List[nib.Nifti1Image], mag: List[nib.Nifti1Image], TEs: Union[List[float], Tuple[float], npt.NDArray[np.float32]], mask: Union[nib.Nifti1Image, SimpleNamespace, None] = None, automask: bool = True, border_size: int = 5, - border_filt: Tuple[int, int] = (1, 5), - svd_filt: int = 10, frames: Union[List[int], None] = None, n_cpus: int = 4, debug: bool = False, wrap_limit: bool = False, -) -> nib.Nifti1Image: - """Unwrap phase of data weighted by magnitude data and compute field maps. This makes a call - to the ROMEO phase unwrapping algorithm for each frame. To learn more about ROMEO, see this paper: - - Dymerska, B., Eckstein, K., Bachrata, B., Siow, B., Trattnig, S., Shmueli, K., Robinson, S.D., 2020. - Phase Unwrapping with a Rapid Opensource Minimum Spanning TreE AlgOrithm (ROMEO). - Magnetic Resonance in Medicine. https://doi.org/10.1002/mrm.28563 - - Parameters - ---------- - phase : List[nib.Nifti1Image] - Phases to unwrap - mag : List[nib.Nifti1Image] - Magnitudes associated with each phase - TEs : Union[List[float], Tuple[float], npt.NDArray[np.float32]] - Echo times associated with each phase (in ms) - mask : nib.Nifti1Image, optional - Boolean mask, by default None - automask : bool, optional - Automatically generate a mask (ignore mask option), by default True - border_size : int, optional - Size of border in automask, by default 5 - border_filt : Tuple[int, int], optional - Number of SVD components for each step of border filtering, by default (1, 5) - svd_filt : int, optional - Number of SVD components to use for filtering of field maps, by default 30 - frames : List[int], optional - Only process these frame indices, by default None (which means all frames) - n_cpus : int, optional - Number of CPUs to use, by default 4 - debug : bool, optional - Debug mode, by default False +) -> Tuple[nib.Nifti1Image, nib.Nifti1Image]: + # make sure affines/shapes are all correct + check_affines(phase, mag) - Returns - ------- - nib.Nifti1Image - Field maps in Hz - """ # check TEs if < 0.1, tell user they probably need to convert to ms if np.min(TEs) < 0.1: logging.warning( - "WARNING: TEs are unusually small. Your inputs may be incorrect. Did you forget to convert to ms?" + "WARNING: TEs are unusually small. Your inputs may be incorrect. " + "Did you forget to convert to ms?" ) # convert TEs to np array TEs = cast(npt.NDArray[np.float32], np.array(TEs)) - # make sure affines/shapes are all correct - for p1, m1 in zip(phase, mag): - for p2, m2 in zip(phase, mag): - if not ( - np.allclose(p1.affine, p2.affine, rtol=1e-3, atol=1e-3) - and np.allclose(p1.shape, p2.shape, rtol=1e-3, atol=1e-3) - and np.allclose(m1.affine, m2.affine, rtol=1e-3, atol=1e-3) - and np.allclose(m1.shape, m2.shape, rtol=1e-3, atol=1e-3) - and np.allclose(p1.affine, m1.affine, rtol=1e-3, atol=1e-3) - and np.allclose(p1.shape, m1.shape, rtol=1e-3, atol=1e-3) - and np.allclose(p2.affine, m2.affine, rtol=1e-3, atol=1e-3) - and np.allclose(p2.shape, m2.shape, rtol=1e-3, atol=1e-3) - ): - raise ValueError("Affines/Shapes of images do not all match.") - - # check if data is 4D or 3D - if len(phase[0].shape) == 3: - # set total number of frames to 1 - n_frames = 1 - # convert data to 4D - phase = [nib.Nifti1Image(p.get_fdata()[..., np.newaxis], p.affine, p.header) for p in phase] - mag = [nib.Nifti1Image(m.get_fdata()[..., np.newaxis], m.affine, m.header) for m in mag] - elif len(phase[0].shape) == 4: - # if frames is None, set it to all frames - if frames is None: - frames = list(range(phase[0].shape[-1])) - # get the total number of frames - n_frames = len(frames) - else: - raise ValueError("Data must be 3D or 4D.") # frames should be a list at this point - frames = cast(List[int], frames) + if frames is None: + frames = list(range(phase[0].shape[3])) + else: + frames = cast(List[int], frames) + n_frames = len(frames) # check echo times = number of mag and phase images if len(TEs) != len(phase) or len(TEs) != len(mag): raise ValueError("Number of echo times must equal number of mag and phase images.") # allocate space for field maps and unwrapped - field_maps = np.zeros((*phase[0].shape[:3], n_frames), dtype=np.float32) unwrapped = np.zeros((*phase[0].shape[:3], len(TEs), n_frames), dtype=np.float32) # array for storing auto-generated masks new_masks = np.zeros((*mag[0].shape[:3], len(frames)), dtype=np.int8) @@ -790,9 +720,66 @@ def post_temporal_consistency_check(idx, result): logging.info("Saving masks..") nib.Nifti1Image(new_masks, phase[0].affine, phase[0].header).to_filename("masks.nii") + return unwrapped, new_masks + + +def compute_field_maps( + unwrapped: npt.NDArray[np.float32], + new_masks: npt.NDArray[np.int8], + img: nib.Nifti1Image, + mag: List[nib.Nifti1Image], + TEs: Union[List[float], Tuple[float], npt.NDArray[np.float32]], + border_filt: Tuple[int, int] = (1, 5), + svd_filt: int = 10, + frames: Union[List[int], None] = None, + n_cpus: int = 4, +) -> nib.Nifti1Image: + """Unwrap phase of data weighted by magnitude data and compute field maps. This makes a call + to the ROMEO phase unwrapping algorithm for each frame. To learn more about ROMEO, see this paper: + + Dymerska, B., Eckstein, K., Bachrata, B., Siow, B., Trattnig, S., Shmueli, K., Robinson, S.D., 2020. + Phase Unwrapping with a Rapid Opensource Minimum Spanning TreE AlgOrithm (ROMEO). + Magnetic Resonance in Medicine. https://doi.org/10.1002/mrm.28563 + + Parameters + ---------- + phase : List[nib.Nifti1Image] + Phases to unwrap + mag : List[nib.Nifti1Image] + Magnitudes associated with each phase + TEs : Union[List[float], Tuple[float], npt.NDArray[np.float32]] + Echo times associated with each phase (in ms) + mask : nib.Nifti1Image, optional + Boolean mask, by default None + automask : bool, optional + Automatically generate a mask (ignore mask option), by default True + border_size : int, optional + Size of border in automask, by default 5 + border_filt : Tuple[int, int], optional + Number of SVD components for each step of border filtering, by default (1, 5) + svd_filt : int, optional + Number of SVD components to use for filtering of field maps, by default 30 + frames : List[int], optional + Only process these frame indices, by default None (which means all frames) + n_cpus : int, optional + Number of CPUs to use, by default 4 + debug : bool, optional + Debug mode, by default False + + Returns + ------- + nib.Nifti1Image + Field maps in Hz + """ + n_frames = img.shape[3] + field_maps = np.zeros((*img.shape[:3], n_frames), dtype=np.float32) + + # convert TEs to np array + TEs = cast(npt.NDArray[np.float32], np.array(TEs)) + # compute field maps on temporally consistent unwrapped phase def field_map_iterator(field_maps, unwrapped, mag, TEs): - logging.info(f"Running field map computation...") + logging.info("Running field map computation...") # convert TEs to a matrix TEs_mat = TEs[:, np.newaxis] for frame_num in range(unwrapped.shape[-1]): @@ -817,11 +804,11 @@ def post_field_map(idx, result): svd_filtering( field_maps, new_masks, - phase[0].header.get_zooms()[0], # type: ignore + img.header.get_zooms()[0], # type: ignore n_frames, border_filt, svd_filt, ) # return the field map as a nifti image - return nib.Nifti1Image(field_maps[..., frames], phase[0].affine, phase[0].header) + return nib.Nifti1Image(field_maps[..., frames], img.affine, img.header) diff --git a/warpkit/utilities.py b/warpkit/utilities.py index 0b6407e..94c551b 100644 --- a/warpkit/utilities.py +++ b/warpkit/utilities.py @@ -759,3 +759,22 @@ def compute_jacobian_determinant(displacement_field: nib.Nifti1Image) -> nib.Nif # return jacobian determinant return cast(nib.Nifti1Image, jacobian_determinant_image) + + +def check_affines(phase, mag): + """Make sure affines/shapes are all correct.""" + for p1, m1 in zip(phase, mag): + for p2, m2 in zip(phase, mag): + if not ( + np.allclose(p1.affine, p2.affine, rtol=1e-3, atol=1e-3) + and np.allclose(m1.affine, m2.affine, rtol=1e-3, atol=1e-3) + and p1.shape == p2.shape + and m1.shape == m2.shape + ): + print(p1.affine, p2.affine) + print(p1.affine - p2.affine) + print(p1.shape, p2.shape) + print(m1.affine, m2.affine) + print(m1.affine - m2.affine) + print(m1.shape, m2.shape) + raise ValueError("Affines and shapes must match")