Add command-line interface for phase unwrapping

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. "

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.")