Merge pull request InsightSoftwareConsortium#4648 from thewtex/master

BUG: array_from_image shape for VectorImage with 1 component
thewtex · May 7, 2024 · 9b8963e · 9b8963e
2 parents e2cdab4 + 2447a81
commit 9b8963e
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git a/Modules/Bridge/NumPy/wrapping/PyBuffer.i.in b/Modules/Bridge/NumPy/wrapping/PyBuffer.i.in
@@ -9,6 +9,7 @@
         i.e. k,j,i versus i,j,k. However C-order indexing is expected by most
         algorithms in NumPy / SciPy.
         """
+        import itk
 
         if image.GetBufferPointer() is None:
             return None
@@ -24,7 +25,7 @@
         dim     = len(itksize)
         shape   = [int(itksize[idx]) for idx in range(dim)]
 
-        if(image.GetNumberOfComponentsPerPixel() > 1):
+        if image.GetNumberOfComponentsPerPixel() > 1 or isinstance(image, itk.VectorImage):
             shape = [image.GetNumberOfComponentsPerPixel(), ] + shape
 
         if keep_axes == False:

diff --git a/Modules/Bridge/NumPy/wrapping/test/itkPyBufferTest.py b/Modules/Bridge/NumPy/wrapping/test/itkPyBufferTest.py
@@ -197,6 +197,19 @@ def test_NumPyBridge_itkVectorImage(self):
         self.assertEqual(vectorndarr.dtype, convertedvectorImage.dtype)
         self.assertTupleEqual(vectorndarr.shape, convertedvectorImage.shape)
 
+        vectorImage = VectorImageType.New()
+        vectorImage.SetRegions(region)
+        vectorImage.SetNumberOfComponentsPerPixel(1)
+        vectorImage.Allocate()
+        vectorndarr = itk.PyBuffer[VectorImageType].GetArrayViewFromImage(vectorImage)
+
+        convertedvectorImage = itk.PyBuffer[VectorImageType].GetImageViewFromArray(
+            vectorndarr, is_vector=True
+        )
+        self.assertEqual(vectorndarr.ndim, convertedvectorImage.ndim)
+        self.assertEqual(vectorndarr.dtype, convertedvectorImage.dtype)
+        self.assertTupleEqual(vectorndarr.shape, convertedvectorImage.shape)
+
     def test_NumPyBridge_itkRGBImage(self):
         "Try to convert an RGB ITK image to NumPy array view"
 

diff --git a/Wrapping/Generators/Python/PyBase/pyBase.i b/Wrapping/Generators/Python/PyBase/pyBase.i
@@ -514,8 +514,10 @@ str = str
             def ndim(self):
                 """Equivalent to the np.ndarray ndim attribute when converted
                 to an image with itk.array_view_from_image."""
+                import itk
+
                 spatial_dims = self.GetImageDimension()
-                if self.GetNumberOfComponentsPerPixel() > 1:
+                if self.GetNumberOfComponentsPerPixel() > 1 or isinstance(self, itk.VectorImage):
                     return spatial_dims + 1
                 else:
                     return spatial_dims
@@ -524,11 +526,13 @@ str = str
             def shape(self):
                 """Equivalent to the np.ndarray shape attribute when converted
                 to an image with itk.array_view_from_image."""
+                import itk
+
                 itksize = self.GetLargestPossibleRegion().GetSize()
                 dim = len(itksize)
                 result = [int(itksize[idx]) for idx in range(dim)]
 
-                if(self.GetNumberOfComponentsPerPixel() > 1):
+                if self.GetNumberOfComponentsPerPixel() > 1 or isinstance(self, itk.VectorImage):
                     result = [self.GetNumberOfComponentsPerPixel(), ] + result
                 # ITK is C-order. The shape needs to be reversed unless we are a view on
                 # a NumPy array that is Fortran-order.