BUG: Keep itk.Image reference in NumPy array views

thewtex · thewtex · commit 87bc23f72353 · 2020-04-22T10:38:17.000-04:00
Prevent the view'ed image from going out of scope by holding a reference in a new class, NDArrayITKBase. The object is kept in the `itk_base` attributed -- the typical `base` attribute NumPy uses for this purpose is not writable. The class is created according to the documentation here: https://numpy.org/doc/stable/user/basics.subclassing.html
diff --git a/Modules/Bridge/NumPy/wrapping/PyBuffer.i.in b/Modules/Bridge/NumPy/wrapping/PyBuffer.i.in
@@ -8,10 +8,6 @@
         indexing. This is the reverse of how indices are specified in ITK,
         i.e. k,j,i versus i,j,k. However C-order indexing is expected by most
         algorithms in NumPy / SciPy.
-
-        Warning: No copy of the data is performed. Using an array
-        view after its source image has been deleted can results in corrupt values
-        or a segfault.
         """
 
         if not HAVE_NUMPY:
@@ -37,11 +33,12 @@
             shape = shape[::-1]
 
         pixelType     = "@PixelType@"
-        numpydatatype = _get_numpy_pixelid(pixelType)
+        numpy_dtype = _get_numpy_pixelid(pixelType)
         memview       = itkPyBuffer@PyBufferTypes@._GetArrayViewFromImage(image)
-        ndarrview  = numpy.asarray(memview).view(dtype = numpydatatype).reshape(shape).view(numpy.ndarray)
+        ndarr_view  = np.asarray(memview).view(dtype = numpy_dtype).reshape(shape).view(np.ndarray)
+        itk_view = NDArrayITKBase(ndarr_view, image)
 
-        return ndarrview
+        return itk_view
 
     GetArrayViewFromImage = staticmethod(GetArrayViewFromImage)
 
@@ -59,7 +56,7 @@
         arrayView = itkPyBuffer@PyBufferTypes@.GetArrayViewFromImage(image, keep_axes, update)
 
         # perform deep copy of the image buffer
-        arr = numpy.array(arrayView, copy=True)
+        arr = np.array(arrayView, copy=True)
 
         return arr
 
@@ -77,10 +74,10 @@
         C-order indexing, i.e. k,j,i, the image Size will have the dimensions
         reversed from the array shape.
 
-        Therefore, since the *numpy.transpose* operator on a 2D array simply
+        Therefore, since the *np.transpose* operator on a 2D array simply
         inverts the indexing scheme, the Image representation will be the
         same for an array and its transpose. If flipping is desired, see
-        *numpy.reshape*.
+        *np.reshape*.
         """
 
         if not HAVE_NUMPY:
@@ -89,7 +86,7 @@
         assert ndarr.ndim in ( 2, 3, 4 ), \
             "Only arrays of 2, 3 or 4 dimensions are supported."
         if not ndarr.flags['C_CONTIGUOUS'] and not ndarr.flags['F_CONTIGUOUS']:
-            ndarr = numpy.ascontiguousarray(ndarr)
+            ndarr = np.ascontiguousarray(ndarr)
 
         if ( ndarr.ndim == 3 and is_vector ) or (ndarr.ndim == 4):
             if( ndarr.flags['C_CONTIGUOUS'] ):
@@ -120,10 +117,10 @@
         C-order indexing, i.e. k,j,i, the image Size will have the dimensions
         reversed from the array shape.
 
-        Therefore, since the *numpy.transpose* operator on a 2D array simply
+        Therefore, since the *np.transpose* operator on a 2D array simply
         inverts the indexing scheme, the Image representation will be the
         same for an array and its transpose. If flipping is desired, see
-        *numpy.reshape*.
+        *np.reshape*.
         """
 
         # Create a temporary image view of the array
diff --git a/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init b/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init
@@ -2,7 +2,19 @@
 
 HAVE_NUMPY = True
 try:
-  import numpy
+  import numpy as np
+  class NDArrayITKBase(np.ndarray):
+      """A numpy array that provides a view on the data associated with an optional itk "base" object."""
+
+      def __new__(cls, input_array, itk_base=None):
+          obj = np.asarray(input_array).view(cls)
+          obj.itk_base = itk_base
+          return obj
+
+      def __array_finalize__(self, obj):
+          if obj is None: return
+          self.itk_base = getattr(obj, 'itk_base', None)
+
 except ImportError:
   HAVE_NUMPY = False
 
@@ -12,21 +24,21 @@ def _get_numpy_pixelid(itk_Image_type):
     if not HAVE_NUMPY:
         raise ImportError('Numpy not available.')
     # This is a Mapping from numpy array types to itk pixel types.
-    _np_itk = {"UC":numpy.uint8,
-               "US":numpy.uint16,
-               "UI":numpy.uint32,
-               "UL":numpy.uint64,
-               "SC":numpy.int8,
-               "SS":numpy.int16,
-               "SI":numpy.int32,
-               "SL":numpy.int64,
-               "F":numpy.float32,
-               "D":numpy.float64,
+    _np_itk = {"UC":np.uint8,
+               "US":np.uint16,
+               "UI":np.uint32,
+               "UL":np.uint64,
+               "SC":np.int8,
+               "SS":np.int16,
+               "SI":np.int32,
+               "SL":np.int64,
+               "F":np.float32,
+               "D":np.float64,
                 }
     import os
     if os.name == 'nt':
-        _np_itk['UL'] = numpy.uint32
-        _np_itk['SL'] = numpy.int32
+        _np_itk['UL'] = np.uint32
+        _np_itk['SL'] = np.int32
     try:
         return _np_itk[itk_Image_type]
     except KeyError as e:
diff --git a/Modules/Bridge/NumPy/wrapping/PyVnlVectorBuffer.i.in b/Modules/Bridge/NumPy/wrapping/PyVnlVectorBuffer.i.in
@@ -2,12 +2,7 @@
     %pythoncode %{
 
     def GetArrayViewFromVnlVector(vnl_vector):
-        """  Get a NumPy array view of a VNL vector.
-
-        Warning: No copy of the data is performed. Using an array
-        view after its source vector has been deleted can results in corrupt values
-        or a segfault.
-        """
+        """Get a NumPy array view of a VNL vector."""
 
         if not HAVE_NUMPY:
             raise ImportError('Numpy not available.')
@@ -16,11 +11,12 @@
         shape   = [itksize]
 
         pixelType     = "@PixelType@"
-        numpydatatype = _get_numpy_pixelid(pixelType)
+        numpy_dtype = _get_numpy_pixelid(pixelType)
         memview       = itkPyVnl@PyVnlTypes@._GetArrayViewFromVnlVector(vnl_vector)
-        ndarrview  = numpy.asarray(memview).view(dtype = numpydatatype).reshape(shape).view(numpy.ndarray)
+        ndarr_view  = np.asarray(memview).view(dtype = numpy_dtype).reshape(shape).view(np.ndarray)
+        itk_view = NDArrayITKBase(ndarr_view, vnl_vector)
 
-        return ndarrview
+        return itk_view
 
     GetArrayViewFromVnlVector = staticmethod(GetArrayViewFromVnlVector)
 
@@ -33,7 +29,7 @@
         arrayView = itkPyVnl@PyVnlTypes@.GetArrayViewFromVnlVector(vnl_vector)
 
         # perform deep copy of the buffer
-        return numpy.array(arrayView, copy=True)
+        return np.array(arrayView, copy=True)
 
     GetArrayFromVnlVector = staticmethod(GetArrayFromVnlVector)
 
@@ -51,20 +47,15 @@
         assert ndarr.ndim == 1 , \
             "Only arrays of 1 dimension are supported."
 
-        ndarr = numpy.ascontiguousarray(ndarr)
+        ndarr = np.ascontiguousarray(ndarr)
         vec = itkPyVnl@PyVnlTypes@._GetVnlVectorFromArray( ndarr, ndarr.shape)
 
         return vec
 
     GetVnlVectorFromArray = staticmethod(GetVnlVectorFromArray)
 
     def GetArrayViewFromVnlMatrix(vnl_matrix):
-        """  Get a NumPy array view of a VNL matrix.
-
-        Warning: No copy of the data is performed. Using an array
-        view after its source matrix has been deleted can results in corrupt values
-        or a segfault.
-        """
+        """Get a NumPy array view of a VNL matrix."""
 
         if not HAVE_NUMPY:
             raise ImportError('Numpy not available.')
@@ -75,11 +66,13 @@
         shape   = [rows,cols]
 
         pixelType     = "@PixelType@"
-        numpydatatype = _get_numpy_pixelid(pixelType)
+        numpy_dtype = _get_numpy_pixelid(pixelType)
         memview       = itkPyVnl@PyVnlTypes@._GetArrayViewFromVnlMatrix(vnl_matrix)
-        ndarrview  = numpy.asarray(memview).view(dtype = numpydatatype).reshape(shape).view(numpy.ndarray)
+        ndarr_view  = np.asarray(memview).view(dtype = numpy_dtype).reshape(shape).view(np.ndarray)
+
+        itk_view = NDArrayITKBase(ndarr_view, vnl_matrix)
 
-        return ndarrview
+        return itk_view
 
     GetArrayViewFromVnlMatrix = staticmethod(GetArrayViewFromVnlMatrix)
 
@@ -92,7 +85,7 @@
         arrayView = itkPyVnl@PyVnlTypes@.GetArrayViewFromVnlMatrix(vnl_matrix)
 
         # perform deep copy of the buffer
-        return numpy.array(arrayView, copy=True)
+        return np.array(arrayView, copy=True)
 
     GetArrayFromVnlMatrix = staticmethod(GetArrayFromVnlMatrix)
 
@@ -111,7 +104,7 @@
         assert ndarr.ndim == 2 , \
             "Only arrays of 2 dimensions are supported."
 
-        ndarr = numpy.ascontiguousarray(ndarr)
+        ndarr = np.ascontiguousarray(ndarr)
         mat = itkPyVnl@PyVnlTypes@._GetVnlMatrixFromArray( ndarr, ndarr.shape)
 
         return mat
diff --git a/Modules/Bridge/NumPy/wrapping/test/itkPyBufferTest.py b/Modules/Bridge/NumPy/wrapping/test/itkPyBufferTest.py
@@ -46,6 +46,9 @@ def test_NumPyBridge_itkScalarImage(self):
         scalarImage.Allocate();
 
         scalarndarr           = itk.PyBuffer[ScalarImageType].GetArrayViewFromImage(scalarImage)
+        self.assertTrue(hasattr(scalarndarr, 'itk_base'))
+        self.assertTrue(scalarndarr.itk_base, scalarImage)
+        self.assertTrue(isinstance(scalarndarr, np.ndarray))
         convertedscalarImage  = itk.PyBuffer[ScalarImageType].GetImageViewFromArray(scalarndarr)
 
         ScalarDiffFilterType  = itk.ComparisonImageFilter[ScalarImageType, ScalarImageType]
