ENH: Handle implicit antsImage -> numpy conversion

Allow images to be put into an array of objects, but not array of scalar types

This prevents obscure errors when an operation is meant to operate on img.numpy() but is passed img instead.

Author: cookpa

ants/core/ants_image.py

@@ -38,6 +38,11 @@
 
 class ANTsImage(object):
 
+    # This makes sure that numpy operations (e.g. np.sum(img)) call our __array__ method so that users get a helpful
+    # error message instead of the numpy fallback to an object array
+    __array_priority__ = 10000
+
+
     def __init__(self, pointer):
         """
         Initialize an ANTsImage.
@@ -571,8 +576,20 @@ def __setitem__(self, idx, value):
     def __iter__(self):
         # Do not allow iteration on ANTsImage. Builtin iteration, eg sum(), will generally be much slower
         # than using numpy methods. We need to explicitly disallow it to prevent breaking object state.
-        raise TypeError("ANTsImage is not iterable. See docs for available functions, or use numpy.")
-
+        raise TypeError("ANTsImage is not iterable. See docs for available functions, or convert to numpy.")
+
+    def __array__(self, dtype=None):
+        if dtype is not None and dtype is np.dtype('O'):
+            # Allow conversion to an object array (eg np.array([img1, img2, img3], dtype=object))
+            out = np.empty((), dtype=object)
+            out[()] = self
+            return out
+        else:
+            # Disallow implicit conversion to numeric numpy array. This prevents using ANTsImage objects in numpy functions
+            # as that's complicated to handle correctly. Users should explicitly convert to numpy array using the .numpy()
+            # function
+            raise TypeError("ANTsImage cannot be implicitly converted to a numeric array. Use the .numpy() method to obtain a "
+                        "copy of the image data as a numpy array. If you want a numpy array of ANTsImage objects, use np.array(..., dtype=object).")
 
     def __repr__(self):
         if self.dimension == 3:

tests/test_plotting.py

@@ -13,8 +13,8 @@
 
 # Prevent displaying figures
 import matplotlib as mpl
-backend_ =  mpl.get_backend() 
-mpl.use("Agg")  
+backend_ =  mpl.get_backend()
+mpl.use("Agg")
 
 import numpy as np
 import numpy.testing as nptest
@@ -39,7 +39,7 @@ def test_plot_example(self):
             ants.plot(img, overlay=img*2)
             ants.plot(img, overlay=img*2)
             ants.plot(img, filename=filename)
-    
+
     def test_extra_plot(self):
         img = ants.image_read(ants.get_ants_data('r16'))
         ants.plot(img, overlay=img*2, domain_image_map=ants.image_read(ants.get_data('r64')))
@@ -48,18 +48,18 @@ def test_extra_plot(self):
         ants.plot(img, crop=True)
 
         img = ants.image_read(ants.get_ants_data('mni'))
-        ants.plot(img, overlay=img*2, 
+        ants.plot(img, overlay=img*2,
                   domain_image_map=ants.image_read(ants.get_data('mni')).resample_image((4,4,4)))
 
         img = ants.image_read(ants.get_ants_data('mni'))
         ants.plot(img, overlay=img*2, reorient=True, crop=True)
-        
+
     def test_random(self):
         img = ants.image_read(ants.get_ants_data('r16'))
         img3 = ants.image_read(ants.get_data('r64'))
         img2 = ants.image_read(ants.get_ants_data('mni'))
         imgv = ants.merge_channels([img2])
-        
+
         ants.plot(img2, axis='x', scale=True, ncol=1)
         ants.plot(img2, axis='y', scale=(0.05, 0.95))
         ants.plot(img2, axis='z', slices=[10,20,30], title='Test', cbar=True,
@@ -76,18 +76,18 @@ def test_random(self):
             ants.plot(ants.get_ants_data('r16'), overlay=ants.merge_channels([img,img]))
         ants.plot(ants.from_numpy(np.zeros((100,100))))
         ants.plot(img.clone('unsigned int'))
-        
+
         ants.plot(img, domain_image_map=img3)
-        
+
         with self.assertRaises(Exception):
             ants.plot(123)
-        
-        
-        
-        
-        
-        
-        
+
+
+
+
+
+
+
 class TestModule_plot_ortho(unittest.TestCase):
 
     def setUp(self):
@@ -102,15 +102,15 @@ def test_plot_example(self):
         for img in self.imgs:
             ants.plot_ortho(img)
             ants.plot_ortho(img, filename=filename)
-            
+
     def test_plot_extra(self):
         img = ants.image_read(ants.get_ants_data('mni'))
-        ants.plot_ortho(img, overlay=img*2, 
+        ants.plot_ortho(img, overlay=img*2,
                   domain_image_map=ants.image_read(ants.get_data('mni')))
 
         img = ants.image_read(ants.get_ants_data('mni'))
         ants.plot_ortho(img, overlay=img*2, reorient=True, crop=True)
-        
+
     def test_random_params(self):
         img = ants.image_read(ants.get_ants_data('mni')).resample_image((4,4,4))
         img2 = ants.image_read(ants.get_data('r16'))
@@ -121,25 +121,25 @@ def test_random_params(self):
             ants.plot_ortho(img2)
         with self.assertRaises(Exception):
             ants.plot_ortho(img, overlay=img2)
-            
+
         imgx = img.clone()
         imgx.set_spacing((3,3,3))
         ants.plot_ortho(img,overlay=imgx)
         ants.plot_ortho(img.clone('unsigned int'),overlay=img, blend=True)
-        
+
         imgx = img.clone()
         imgx.set_spacing((10,1,1))
         ants.plot_ortho(imgx)
-        
+
         ants.plot_ortho(img, flat=True, title='Test', text='This is a test')
         ants.plot_ortho(img, title='Test', text='This is a test', cbar=True)
-        
+
         with self.assertRaises(Exception):
             ants.plot_ortho(img, domain_image_map=123)
-            
+
         with self.assertRaises(Exception):
             ants.plot_orto(ants.merge_channels([img,img]))
-        
+
 
 class TestModule_plot_ortho_stack(unittest.TestCase):
 
@@ -153,7 +153,7 @@ def test_plot_example(self):
         filename = mktemp(suffix='.png')
         ants.plot_ortho_stack([self.img, self.img])
         ants.plot_ortho_stack([self.img, self.img], filename=filename)
-        
+
     def test_extra_ortho_stack(self):
         img = ants.image_read(ants.get_ants_data('mni'))
         ants.plot_ortho_stack([img, img], overlays=[img*2, img*2],
@@ -185,9 +185,9 @@ def setUp(self):
         mni3 = mni1.smooth_image(2.)
         mni4 = mni1.smooth_image(3.)
         self.images3d = np.asarray([[mni1, mni2],
-                                    [mni3, mni4]])
+                                    [mni3, mni4]], dtype='object')
         self.images2d = np.asarray([[mni1.slice_image(2,100), mni2.slice_image(2,100)],
-                                    [mni3.slice_image(2,100), mni4.slice_image(2,100)]])
+                                    [mni3.slice_image(2,100), mni4.slice_image(2,100)]],dtype='object')
 
     def tearDown(self):
         pass
@@ -198,19 +198,19 @@ def test_plot_example(self):
         ants.plot_grid(self.images3d)
         # should work with 2d images
         ants.plot_grid(self.images2d)
-        
+
     def test_examples(self):
         mni1 = ants.image_read(ants.get_data('mni'))
         mni2 = mni1.smooth_image(1.)
         mni3 = mni1.smooth_image(2.)
         mni4 = mni1.smooth_image(3.)
         images = np.asarray([[mni1, mni2],
-                            [mni3, mni4]])
+                            [mni3, mni4]], dtype='object')
         slices = np.asarray([[100, 100],
-                            [100, 100]])
+                            [100, 100]], dtype='object')
         ants.plot_grid(images=images, slices=slices, title='2x2 Grid')
         images2d = np.asarray([[mni1.slice_image(2,100), mni2.slice_image(2,100)],
-                            [mni3.slice_image(2,100), mni4.slice_image(2,100)]])
+                            [mni3.slice_image(2,100), mni4.slice_image(2,100)]], dtype='object')
         ants.plot_grid(images=images2d, title='2x2 Grid Pre-Sliced')
         ants.plot_grid(images.reshape(1,4), slices.reshape(1,4), title='1x4 Grid')
         ants.plot_grid(images.reshape(4,1), slices.reshape(4,1), title='4x1 Grid')
@@ -228,11 +228,11 @@ def test_examples(self):
 
         # Making a publication-quality image
         images = np.asarray([[mni1, mni2, mni2],
-                            [mni3, mni4, mni4]])
+                            [mni3, mni4, mni4]], dtype='object')
         slices = np.asarray([[100, 100, 100],
-                            [100, 100, 100]])
+                            [100, 100, 100]], dtype='object')
         axes = np.asarray([[0, 1, 2],
-                        [0, 1, 2]])
+                        [0, 1, 2]], dtype='object')
         ants.plot_grid(images, slices, axes, title='Publication Figures with ANTsPy',
                     tfontsize=20, title_dy=0.03, title_dx=-0.04,
                     rlabels=['Row 1', 'Row 2'],
@@ -246,15 +246,15 @@ def setUp(self):
         pass
     def tearDown(self):
         pass
-    
+
     def test_random_ortho_stack_params(self):
         img = ants.image_read(ants.get_data('mni')).resample_image((4,4,4))
         img2 = ants.image_read(ants.get_data('r16')).resample_image((4,4))
-        
+
         ants.plot_ortho_stack([ants.get_data('mni'), ants.get_data('mni')])
         ants.plot_ortho_stack([ants.get_data('mni'), ants.get_data('mni')],
                               overlays=[ants.get_data('mni'), ants.get_data('mni')])
-        
+
         with self.assertRaises(Exception):
             ants.plot_ortho_stack([1,2,3])
         with self.assertRaises(Exception):
@@ -263,18 +263,18 @@ def test_random_ortho_stack_params(self):
             ants.plot_ortho_stack([img,img], overlays=[img2,img2])
         with self.assertRaises(Exception):
             ants.plot_ortho_stack([img,img], overlays=[1,2])
-            
+
         imgx = img.clone()
         imgx.set_spacing((2,2,2))
         ants.plot_ortho_stack([img,imgx])
-        
+
         imgx.set_spacing((2,1,1))
         ants.plot_ortho_stack([imgx,img])
-        
+
         ants.plot_ortho_stack([img,img], scale=True, transpose=True,
                               title='Test', colpad=1, rowpad=1,
                               xyz_lines=True)
         ants.plot_ortho_stack([img,img], scale=(0.05,0.95))
-            
+
 if __name__ == '__main__':
     run_tests()