ENH: Add nib-roi command to crop (maybe flip) axes

nibabel/cmdline/roi.py

@@ -0,0 +1,88 @@
+import sys
+import os
+import argparse
+import nibabel as nb
+
+
+def lossless_slice(img, slicers):
+    if not nb.imageclasses.spatial_axes_first(img):
+        raise ValueError("Cannot slice an image that is not known to have spatial axes first")
+
+    scaling = hasattr(img.header, 'set_slope_inter')
+
+    data = img.dataobj._get_unscaled(slicers) if scaling else img.dataobj[slicers]
+    roi_img = img.__class__(data, affine=img.slicer.slice_affine(slicers), header=img.header)
+
+    if scaling:
+        roi_img.header.set_slope_inter(img.dataobj.slope, img.dataobj.inter)
+    return roi_img
+
+
+def parse_slice(crop, allow_step=True):
+    if crop is None:
+        return slice(None)
+    start, stop, *extra = [int(val) if val else None for val in crop.split(":")]
+    if len(extra) > 1:
+        raise ValueError(f"Cannot parse specification: {crop}")
+    if not allow_step and extra and extra[0] not in (1, None):
+        raise ValueError(f"Step entry not permitted: {crop}")
+
+    step = extra[0] if extra else None
+    if step not in (1, -1, None):
+        raise ValueError(f"Downsampling is not supported: {crop}")
+
+    return slice(start, stop, step)
+
+
+def sanitize(args):
+    # Argparse likes to treat "-1:..." as a flag
+    return [f' {arg}' if arg[0] == '-' and ":" in arg else arg
+            for arg in args]
+
+
+def main(args=None):
+    if args is None:
+        args = sys.argv[1:]
+    parser = argparse.ArgumentParser(
+        description="Crop images to a region of interest",
+        epilog="If a start or stop value is omitted, the start or end of the axis is assumed.")
+    parser.add_argument('--version', action='version', version=nb.__version__)
+    parser.add_argument("-i", metavar="I1:I2[:-1]",
+                        help="Start/stop [flip] along first axis (0-indexed)")
+    parser.add_argument("-j", metavar="J1:J2[:-1]",
+                        help="Start/stop [flip] along second axis (0-indexed)")
+    parser.add_argument("-k", metavar="K1:K2[:-1]",
+                        help="Start/stop [flip] along third axis (0-indexed)")
+    parser.add_argument("-t", metavar="T1:T2", help="Start/stop along fourth axis (0-indexed)")
+    parser.add_argument("in_file", help="Image file to crop")
+    parser.add_argument("out_file", help="Output file name")
+
+    opts = parser.parse_args(args=sanitize(args))
+
+    try:
+        islice = parse_slice(opts.i)
+        jslice = parse_slice(opts.j)
+        kslice = parse_slice(opts.k)
+        tslice = parse_slice(opts.t, allow_step=False)
+    except ValueError as err:
+        print(f"Could not parse input arguments. Reason follows.\n{err}")
+        return 1
+
+    kwargs = {}
+    if os.path.realpath(opts.in_file) == os.path.realpath(opts.out_file):
+        kwargs['mmap'] = False
+    img = nb.load(opts.in_file, **kwargs)
+
+    slicers = (islice, jslice, kslice, tslice)[:img.ndim]
+    expected_shape = nb.fileslice.predict_shape(slicers, img.shape)
+    if any(dim == 0 for dim in expected_shape):
+        print(f"Cannot take zero-length slices. Predicted shape {expected_shape}.")
+        return 1
+
+    try:
+        sliced_img = lossless_slice(img, slicers)
+    except Exception:
+        print("Could not slice image. Full traceback follows.")
+        raise
+    nb.save(sliced_img, opts.out_file)
+    return 0

nibabel/cmdline/tests/test_roi.py

@@ -0,0 +1,152 @@
+import os
+import numpy as np
+import nibabel as nb
+from nibabel.cmdline.roi import lossless_slice, parse_slice, main
+from nibabel.testing import data_path
+
+import unittest
+from unittest import mock
+import pytest
+
+
+def test_parse_slice():
+    assert parse_slice(None) == slice(None)
+    assert parse_slice("1:5") == slice(1, 5)
+    assert parse_slice("1:") == slice(1, None)
+    assert parse_slice(":5") == slice(None, 5)
+    assert parse_slice(":-1") == slice(None, -1)
+    assert parse_slice("-5:-1") == slice(-5, -1)
+    assert parse_slice("1:5:") == slice(1, 5, None)
+    assert parse_slice("1::") == slice(1, None, None)
+    assert parse_slice(":5:") == slice(None, 5, None)
+    assert parse_slice(":-1:") == slice(None, -1, None)
+    assert parse_slice("-5:-1:") == slice(-5, -1, None)
+    assert parse_slice("1:5:1") == slice(1, 5, 1)
+    assert parse_slice("1::1") == slice(1, None, 1)
+    assert parse_slice(":5:1") == slice(None, 5, 1)
+    assert parse_slice(":-1:1") == slice(None, -1, 1)
+    assert parse_slice("-5:-1:1") == slice(-5, -1, 1)
+    assert parse_slice("5:1:-1") == slice(5, 1, -1)
+    assert parse_slice(":1:-1") == slice(None, 1, -1)
+    assert parse_slice("5::-1") == slice(5, None, -1)
+    assert parse_slice("-1::-1") == slice(-1, None, -1)
+    assert parse_slice("-1:-5:-1") == slice(-1, -5, -1)
+
+    # Max of start:stop:step
+    with pytest.raises(ValueError):
+        parse_slice("1:2:3:4")
+    # Integers only
+    with pytest.raises(ValueError):
+        parse_slice("abc:2:3")
+    with pytest.raises(ValueError):
+        parse_slice("1.2:2:3")
+    # Unit steps only
+    with pytest.raises(ValueError):
+        parse_slice("1:5:2")
+
+
+def test_parse_slice_disallow_step():
+    # Permit steps of 1
+    assert parse_slice("1:5", False) == slice(1, 5)
+    assert parse_slice("1:5:", False) == slice(1, 5)
+    assert parse_slice("1:5:1", False) == slice(1, 5, 1)
+    # Disable other steps
+    with pytest.raises(ValueError):
+        parse_slice("1:5:-1", False)
+    with pytest.raises(ValueError):
+        parse_slice("1:5:-2", False)
+
+
+def test_lossless_slice_unknown_axes():
+    img = nb.load(os.path.join(data_path, 'minc1_4d.mnc'))
+    with pytest.raises(ValueError):
+        lossless_slice(img, (slice(None), slice(None), slice(None)))
+
+
+def test_lossless_slice_scaling(tmp_path):
+    fname = tmp_path / 'image.nii'
+    img = nb.Nifti1Image(np.random.uniform(-20000, 20000, (5, 5, 5, 5)), affine=np.eye(4))
+    img.header.set_data_dtype("int16")
+    img.to_filename(fname)
+    img1 = nb.load(fname)
+    sliced_fname = tmp_path / 'sliced.nii'
+    lossless_slice(img1, (slice(None), slice(None), slice(2, 4))).to_filename(sliced_fname)
+    img2 = nb.load(sliced_fname)
+
+    assert np.array_equal(img1.get_fdata()[:, :, 2:4], img2.get_fdata())
+    assert np.array_equal(img1.dataobj.get_unscaled()[:, :, 2:4], img2.dataobj.get_unscaled())
+    assert img1.dataobj.slope == img2.dataobj.slope
+    assert img1.dataobj.inter == img2.dataobj.inter
+
+
+def test_lossless_slice_noscaling(tmp_path):
+    fname = tmp_path / 'image.mgh'
+    img = nb.MGHImage(np.random.uniform(-20000, 20000, (5, 5, 5, 5)).astype("float32"),
+                      affine=np.eye(4))
+    img.to_filename(fname)
+    img1 = nb.load(fname)
+    sliced_fname = tmp_path / 'sliced.mgh'
+    lossless_slice(img1, (slice(None), slice(None), slice(2, 4))).to_filename(sliced_fname)
+    img2 = nb.load(sliced_fname)
+
+    assert np.array_equal(img1.get_fdata()[:, :, 2:4], img2.get_fdata())
+    assert np.array_equal(img1.dataobj.get_unscaled()[:, :, 2:4], img2.dataobj.get_unscaled())
+    assert img1.dataobj.slope == img2.dataobj.slope
+    assert img1.dataobj.inter == img2.dataobj.inter
+
+
+@pytest.mark.parametrize("inplace", (True, False))
+def test_nib_roi(tmp_path, inplace):
+    in_file = os.path.join(data_path, 'functional.nii')
+    out_file = str(tmp_path / 'sliced.nii')
+    in_img = nb.load(in_file)
+
+    if inplace:
+        in_img.to_filename(out_file)
+        in_file = out_file
+
+    retval = main([in_file, out_file, '-i', '1:-1', '-j', '-1:1:-1', '-k', '::', '-t', ':5'])
+    assert retval == 0
+
+    out_img = nb.load(out_file)
+    in_data = in_img.dataobj[:]
+    in_sliced = in_img.slicer[1:-1, -1:1:-1, :, :5]
+    assert out_img.shape == in_sliced.shape
+    assert np.array_equal(in_data[1:-1, -1:1:-1, :, :5], out_img.dataobj)
+    assert np.allclose(in_sliced.dataobj, out_img.dataobj)
+    assert np.allclose(in_sliced.affine, out_img.affine)
+
+
+@pytest.mark.parametrize("args, errmsg", (
+    (("-i", "1:1"), "Cannot take zero-length slice"),
+    (("-j", "1::2"), "Downsampling is not supported"),
+    (("-t", "5::-1"), "Step entry not permitted"),
+))
+def test_nib_roi_bad_slices(capsys, args, errmsg):
+    in_file = os.path.join(data_path, 'functional.nii')
+
+    retval = main([in_file, os.devnull, *args])
+    assert retval != 0
+    captured = capsys.readouterr()
+    assert errmsg in captured.out
+
+
+def test_entrypoint(capsys):
+    # Check that we handle missing args as expected
+    with mock.patch("sys.argv", ["nib-roi", "--help"]):
+        try:
+            retval = main()
+        except SystemExit:
+            pass
+        else:
+            assert False, "argparse exits on --help. If changing to another parser, update test."
+    captured = capsys.readouterr()
+    assert captured.out.startswith("usage: nib-roi")
+
+
+def test_nib_roi_unknown_axes(capsys):
+    in_file = os.path.join(data_path, 'minc1_4d.mnc')
+    with pytest.raises(ValueError):
+        main([in_file, os.devnull, "-i", ":"])
+    captured = capsys.readouterr()
+    assert "Could not slice image." in captured.out

setup.cfg

@@ -77,6 +77,7 @@ console_scripts =
     nib-nifti-dx=nibabel.cmdline.nifti_dx:main
     nib-tck2trk=nibabel.cmdline.tck2trk:main
     nib-trk2tck=nibabel.cmdline.trk2tck:main
+    nib-roi=nibabel.cmdline.roi:main
     parrec2nii=nibabel.cmdline.parrec2nii:main
 
 [options.package_data]