Add an 'axis' parameter to concat_images, plus two tests.

Author: Ben Cipollini

nibabel/funcs.py

@@ -88,8 +88,8 @@ def squeeze_image(img):
                  img.extra)
 
 
-def concat_images(images, check_affines=True):
-    ''' Concatenate images in list to single image, along last dimension
+def concat_images(images, check_affines=True, axis=None):
+    ''' Concatenate images in list to single image, along specified dimension
 
     Parameters
     ----------
@@ -98,7 +98,9 @@ def concat_images(images, check_affines=True):
     check_affines : {True, False}, optional
        If True, then check that all the affines for `images` are nearly
        the same, raising a ``ValueError`` otherwise.  Default is True
-
+    axis : int, optional
+        If None, concatenates on the last dimension.
+        If not None, concatenates on the specified dimension.
     Returns
     -------
     concat_img : ``SpatialImage``
@@ -122,8 +124,13 @@ def concat_images(images, check_affines=True):
         if check_affines:
             if not np.all(img.affine == affine):
                 raise ValueError('Affines do not match')
-        out_data[i] = img.get_data()
-    out_data = np.rollaxis(out_data, 0, len(i0shape)+1)
+        out_data[i] = img.get_data().copy()
+    if axis is not None:
+        out_data = np.concatenate(out_data, axis=axis)
+    elif np.all([d.shape[-1] == 1 for d in out_data]):
+        out_data = np.concatenate(out_data, axis=d.ndim-1)
+    else:
+        out_data = np.rollaxis(out_data, 0, len(i0shape)+1)
     klass = img0.__class__
     return klass(out_data, affine, header)
 

nibabel/tests/test_funcs.py

@@ -59,6 +59,23 @@ def test_concat():
         for img in imgs:
             del(img)
 
+    # Test axis parameter and trailing unary dimension
+    shape_4D = np.asarray(shape + (1,))
+    data0 = np.arange(10).reshape(shape_4D)
+    affine = np.eye(4)
+    img0_mem = Nifti1Image(data0, affine)
+    img1_mem = Nifti1Image(data0 - 10, affine)
+
+    concat_img1 = concat_images([img0_mem, img1_mem])
+    expected_shape1 = shape_4D.copy()
+    expected_shape1[-1] *= 2
+    assert_array_equal(concat_img1.shape, expected_shape1)
+
+    concat_img2 = concat_images([img0_mem, img1_mem], axis=0)
+    expected_shape2 = shape_4D.copy()
+    expected_shape2[0] *= 2
+    assert_array_equal(concat_img2.shape, expected_shape2)
+
 
 def test_closest_canonical():
     arr = np.arange(24).reshape((2,3,4,1))