BUG: Preserve component layout when reconstructing and serializing images

ants/core/ants_image.py

@@ -595,11 +595,23 @@ def __getstate__(self):
         img += 10
         print(img.mean(), img3.mean())
         """
-        return self.numpy(), self.origin, self.spacing, self.direction, self.has_components, self.is_rgb
+        data = self.numpy()
+        if self.has_components and self.channels_first:
+            data = np.moveaxis(data, 0, -1)
+        state = (data, self.origin, self.spacing, self.direction, self.has_components, self.is_rgb)
+        # Keep the six-item state for existing images and retain the component
+        # presentation only when it differs from the default.
+        if self.has_components and self.channels_first:
+            return state + (True,)
+        return state
 
     def __setstate__(self, state):
-        data, origin, spacing, direction, has_components, is_rgb = state
+        if len(state) not in (6, 7):
+            raise ValueError('ANTsImage pickle state must contain 6 or 7 items')
+        data, origin, spacing, direction, has_components, is_rgb = state[:6]
+        channels_first = bool(state[6]) if len(state) == 7 else False
         image = ants.from_numpy(np.copy(data), origin=origin, spacing=spacing, direction=direction, has_components=has_components, is_rgb=is_rgb)
+        image.channels_first = channels_first if has_components else False
         self.__dict__ = image.__dict__
 
 def copy_image_info(reference, target):

ants/core/ants_image_io.py

@@ -389,6 +389,8 @@ def image_read(filename, dimension=None, pixeltype="float", reorient=False):
                 has_components=img_header.get("components", 1) > 1,
                 pixeltype=pixeltype
             )
+            if ants_image.has_components:
+                ants_image.channels_first = bool(img_header.get("channels_first", False))
         else:
             img_header = {}
             ants_image = from_numpy(img_array, pixeltype=pixeltype) # Still does type conversion if pixeltype=None
@@ -521,12 +523,16 @@ def image_write(image, filename, ri=False):
     """
     if filename.endswith(".npy"):
         img_array = image.numpy()
+        if image.has_components and image.channels_first:
+            img_array = np.moveaxis(img_array, 0, -1)
         img_header = {
             "origin": image.origin,
             "spacing": image.spacing,
             "direction": image.direction.tolist(),
             "components": image.components,
         }
+        if image.has_components and image.channels_first:
+            img_header["channels_first"] = True
 
         np.save(filename, img_array)
         with open(filename[:-4] + ".json", "w") as outfile:
@@ -595,25 +601,30 @@ def new_image_like(image, data):
     ---------
     data : ndarray or py::capsule
         New array or pointer for the image.
-        It must have the same shape as the current
-        image data.
+        It must have the same shape and component-axis order as
+        ``image.numpy()``.
 
     Returns
     -------
     ANTsImage
     """
     if not isinstance(data, np.ndarray):
         raise ValueError('data must be a numpy array')
-    if not image.has_components:
-        if data.shape != image.shape:
-            raise ValueError('given array shape (%s) and image array shape (%s) do not match' % (data.shape, image.shape))
-    else:
-        if (data.shape[-1] != image.components) or (data.shape[:-1] != image.shape):
-            raise ValueError('given array shape (%s) and image array shape (%s) do not match' % (data.shape[1:], image.shape))
-
-    return from_numpy(data, origin=image.origin,
+    expected_shape = image.shape
+    if image.has_components:
+        expected_shape = ((image.components,) + image.shape if image.channels_first
+                          else image.shape + (image.components,))
+    if data.shape != expected_shape:
+        raise ValueError('given array shape %s does not match image.numpy() shape %s' %
+                         (data.shape, expected_shape))
+    if image.has_components and image.channels_first:
+        data = np.moveaxis(data, 0, -1)
+
+    new_image = from_numpy(data, origin=image.origin,
         spacing=image.spacing, direction=image.direction,
         has_components=image.has_components)
+    new_image.channels_first = image.channels_first
+    return new_image
 
 def from_numpy_like(data, image):
     return new_image_like(image, data)

ants/utils/nibabel_nifti_to_ants.py

@@ -148,6 +148,8 @@ def to_nibabel_nifti( ants_image, header=None):
 
     if ants_image.has_components:
         # nibabel expects components in the 5th dimension, so array has to have dims (x,y,z,1,components)
+        if ants_image.channels_first:
+            data = np.moveaxis(data, 0, -1)
         data = np.expand_dims(data, axis=3)
 
     spacing = ants_image.spacing

tests/test_core_ants_image.py

@@ -697,5 +697,80 @@ def test_rgb_image(self):
         self.assertTrue(rgb_img.dtype == 'uint8')
 
 
+class TestComponentLayout(unittest.TestCase):
+
+    def setUp(self):
+        self.images = []
+        for shape in ((3, 4), (3, 4, 5)):
+            dimension = len(shape)
+            array = np.arange(np.prod(shape), dtype=np.float32).reshape(shape)
+            scalar = ants.from_numpy(array, origin=tuple(range(dimension)),
+                                     spacing=tuple(range(2, dimension + 2)),
+                                     direction=np.flipud(np.eye(dimension)))
+            for channels_first in (False, True):
+                self.images.append(ants.merge_channels(
+                    [scalar, scalar + 5], channels_first=channels_first))
+
+    def assert_same_geometry(self, image, result):
+        self.assertEqual(result.shape, image.shape)
+        self.assertEqual(result.components, image.components)
+        self.assertEqual(result.origin, image.origin)
+        self.assertEqual(result.spacing, image.spacing)
+        nptest.assert_array_equal(result.direction, image.direction)
+
+    def test_new_image_like_component_layout(self):
+        for image in self.images:
+            for operation in (lambda img: img.new_image_like(img.numpy() * 2),
+                              lambda img: img * 2,
+                              lambda img: img + img):
+                with self.subTest(shape=image.shape, channels_first=image.channels_first,
+                                  operation=operation):
+                    result = operation(image)
+                    self.assert_same_geometry(image, result)
+                    self.assertEqual(result.channels_first, image.channels_first)
+                    nptest.assert_array_equal(result.numpy(), image.numpy() * 2)
+
+    def test_new_image_like_invalid_component_layout(self):
+        for image in self.images:
+            # Supply the opposite component-axis order to check that the error
+            # reports the original shape, not the internally moved shape.
+            wrong_shape = (image.shape + (image.components,) if image.channels_first
+                           else (image.components,) + image.shape)
+            data = np.zeros(wrong_shape, dtype=np.float32)
+            with self.subTest(shape=image.shape, channels_first=image.channels_first):
+                with self.assertRaises(ValueError) as raised:
+                    image.new_image_like(data)
+                self.assertIn(str(wrong_shape), str(raised.exception))
+                self.assertIn(str(image.numpy().shape), str(raised.exception))
+
+    def test_pickle_component_layout(self):
+        import pickle
+        from copy import deepcopy
+        for image in self.images:
+            expected_length = 7 if image.channels_first else 6
+            self.assertEqual(len(image.__getstate__()), expected_length)
+            for operation in (lambda img: pickle.loads(pickle.dumps(img)), deepcopy):
+                with self.subTest(shape=image.shape, channels_first=image.channels_first,
+                                  operation=operation):
+                    result = operation(image)
+                    self.assert_same_geometry(image, result)
+                    self.assertEqual(result.channels_first, image.channels_first)
+                    nptest.assert_array_equal(result.numpy(), image.numpy())
+
+    def test_legacy_six_item_pickle_state(self):
+        for image in self.images:
+            with self.subTest(shape=image.shape, channels_first=image.channels_first):
+                # Prior versions always stored components last in six items.
+                legacy_state = image.__getstate__()[:6]
+                restored = ants.ANTsImage.__new__(ants.ANTsImage)
+                restored.__setstate__(legacy_state)
+                self.assert_same_geometry(image, restored)
+                self.assertFalse(restored.channels_first)
+                expected = image.numpy()
+                if image.channels_first:
+                    expected = np.moveaxis(expected, 0, -1)
+                nptest.assert_array_equal(restored.numpy(), expected)
+
+
 if __name__ == '__main__':
     run_tests()

tests/test_core_ants_image_io.py

@@ -431,6 +431,32 @@ def setUp(self):
         self.header = dict(origin=(2.5, -1.0, 8.0), spacing=(0.5, 1.25, 2.0),
                            direction=[[0, -1, 0], [1, 0, 0], [0, 0, 1]])
 
+    def test_npy_component_layout(self):
+        for array in (self.array[0], self.array):
+            dimension = array.ndim
+            data = np.stack((array, array + 5), axis=-1)
+            image = ants.from_numpy(data, has_components=True,
+                                    origin=self.header['origin'][:dimension],
+                                    spacing=self.header['spacing'][:dimension],
+                                    direction=np.asarray(self.header['direction'])[:dimension, :dimension])
+            for channels_first in (False, True):
+                image.channels_first = channels_first
+                with self.subTest(shape=array.shape, channels_first=channels_first), TemporaryDirectory() as tmp:
+                    filename = os.path.join(tmp, 'image.npy')
+                    ants.image_write(image, filename)
+                    nptest.assert_array_equal(np.load(filename), data)
+                    with open(filename[:-4] + '.json') as stream:
+                        sidecar = json.load(stream)
+                    self.assertEqual(sidecar.get('channels_first', False), channels_first)
+                    result = ants.image_read(filename)
+                    self.assertEqual(result.shape, image.shape)
+                    self.assertEqual(result.components, image.components)
+                    self.assertEqual(result.origin, image.origin)
+                    self.assertEqual(result.spacing, image.spacing)
+                    nptest.assert_array_equal(result.direction, image.direction)
+                    nptest.assert_array_equal(result.numpy(), image.numpy())
+                    self.assertEqual(result.channels_first, channels_first)
+
     def test_npy_sidecar_round_trip(self):
         for name in ('image.npy', 'image.npy.copy.npy',
                      'directory.npy/image.npy', 'directory.npy/image.npy.copy.npy'):

tests/test_utils.py

@@ -1691,6 +1691,23 @@ def test_multicomponent(self):
         self.assertTrue(np.allclose(nib_img_vec.get_fdata(), ants_img_back.get_fdata(), atol=1e-5))
         self.assertTrue(np.allclose(nib_img_vec.affine, ants_img_back.affine, atol=1e-5))
 
+    def test_multicomponent_channels_first(self):
+        scalar_data = np.arange(4 * 5 * 6, dtype=np.float32).reshape(4, 5, 6)
+        scalar = ants.from_numpy(scalar_data)
+        expected = np.stack((scalar_data, scalar_data + 10, scalar_data * 2), axis=-1)
+        for channels_first in (False, True):
+            with self.subTest(channels_first=channels_first):
+                image = ants.merge_channels(
+                    [scalar, scalar + 10, scalar * 2], channels_first=channels_first)
+                nifti = ants.to_nibabel_nifti(image)
+                self.assertEqual(nifti.shape, (4, 5, 6, 1, 3))
+                nptest.assert_array_equal(np.asanyarray(nifti.dataobj), expected[:, :, :, None, :])
+
+                restored = ants.from_nibabel_nifti(nifti)
+                self.assertEqual(restored.shape, image.shape)
+                self.assertEqual(restored.components, image.components)
+                nptest.assert_array_equal(restored.numpy(), expected)
+
 
 if __name__ == "__main__":
     run_tests()