Add support for other dimensions than 3

RTKConsortium · Sep 3, 2024 · 6fe09e2 · 6fe09e2
1 parent b6ac52f
commit 6fe09e2
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/wrapping/CudaImage.i.in b/wrapping/CudaImage.i.in
@@ -4,8 +4,12 @@
     def __cuda_array_interface__(self):
         _pixelType = "@PixelType@"
         _typestr = _get_type_string(_pixelType)
+
+        _itksize = self.GetBufferedRegion().GetSize()
+        _dim     = len(_itksize)
+        _shape   = tuple(int(_itksize[idx]) for idx in range(_dim))
         return {
-            'shape': (self.GetLargestPossibleRegion().GetSize()[0], self.GetLargestPossibleRegion().GetSize()[1], self.GetLargestPossibleRegion().GetSize()[2]),
+            'shape': _shape,
             'data': (int(self.GetCudaDataManager().GetGPUBufferPointer()), False),
             'typestr': _typestr,
             'descr': [('', _typestr)],