Skip to content

Commit 5ee0652

Browse files
committed
WIP: Implement PEP 688 for image and image container
Experiment with implementing PEP 688. Replace __array__ with __buffer__. Convert memoryview to correctly specified n-d view. PEP 688 automatically maintains a a reference to the provider of the buffer. Add incomplete wrapping for the import image container, and buffer interface.
1 parent a9e3442 commit 5ee0652

File tree

6 files changed

+170
-9
lines changed

6 files changed

+170
-9
lines changed

Modules/Bridge/NumPy/include/itkPyBuffer.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,12 @@ class PyBuffer
7272
static PyObject *
7373
_GetArrayViewFromImage(ImageType * image);
7474

75+
/**
76+
* Get an 1-D byte MemoryView of the container's buffer
77+
*/
78+
static PyObject *
79+
_GetMemoryViewFromImportImageContainer(typename ImageType::PixelContainer * container);
80+
7581
/**
7682
* Get an ITK image from a contiguous Python array. Internal helper function for the implementation of
7783
* `itkPyBuffer.GetImageViewFromArray`.

Modules/Bridge/NumPy/include/itkPyBuffer.hxx

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,38 @@ PyBuffer<TImage>::_GetArrayViewFromImage(ImageType * image)
5050
return PyMemoryView_FromBuffer(&pyBuffer);
5151
}
5252

53+
template <class TImage>
54+
PyObject *
55+
PyBuffer<TImage>::_GetMemoryViewFromImportImageContainer(typename ImageType::PixelContainer * container)
56+
{
57+
using ContainerType = typename ImageType::PixelContainer;
58+
Py_buffer pyBuffer{};
59+
60+
if (!container)
61+
{
62+
throw std::runtime_error("Input container is null");
63+
}
64+
65+
void * const buffer = container->GetBufferPointer();
66+
67+
if (!buffer)
68+
{
69+
throw std::runtime_error("Container buffer pointer is null");
70+
}
71+
72+
// If the container does not own the buffer then issue a warning
73+
if (!container->GetContainerManageMemory())
74+
{
75+
PyErr_WarnEx(PyExc_RuntimeWarning, "The ImportImageContainer does not own the exported buffer.", 1);
76+
}
77+
78+
const SizeValueType size = container->Size();
79+
const auto len = static_cast<Py_ssize_t>(size * sizeof(typename ContainerType::Element));
80+
81+
PyBuffer_FillInfo(&pyBuffer, nullptr, buffer, len, 0, PyBUF_CONTIG);
82+
return PyMemoryView_FromBuffer(&pyBuffer);
83+
}
84+
5385
template <class TImage>
5486
auto
5587
PyBuffer<TImage>::_get_image_view_from_contiguous_array(PyObject * arr, PyObject * shape, PyObject * numOfComponent)

Modules/Bridge/NumPy/wrapping/PyBuffer.i.init

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,53 @@ else:
2929
loads = dask_deserialize.dispatch(np.ndarray)
3030
return NDArrayITKBase(loads(header, frames))
3131

32+
def _get_formatstring(itk_Image_type) -> str:
33+
"""Returns the struct format string for a given ITK image type.
34+
35+
Format characters from Python's struct module:
36+
- 'b': signed char (int8)
37+
- 'B': unsigned char (uint8)
38+
- 'h': short (int16)
39+
- 'H': unsigned short (uint16)
40+
- 'i': int (int32)
41+
- 'I': unsigned int (uint32)
42+
- 'l': long (platform dependent)
43+
- 'L': unsigned long (platform dependent)
44+
- 'q': long long (int64)
45+
- 'Q': unsigned long long (uint64)
46+
- 'f': float (float32)
47+
- 'd': double (float64)
48+
"""
49+
50+
# Mapping from ITK pixel type codes to struct format strings
51+
_format_map = {
52+
"UC": "B", # unsigned char
53+
"US": "H", # unsigned short
54+
"UI": "I", # unsigned int
55+
"UL": "L", # unsigned long
56+
"ULL": "Q", # unsigned long long
57+
"SC": "b", # signed char
58+
"SS": "h", # signed short
59+
"SI": "i", # signed int
60+
"SL": "l", # signed long
61+
"SLL": "q", # signed long long
62+
"F": "f", # float
63+
"D": "d", # double
64+
"PF2": "f", # Point<float, 2> - use float for components
65+
"PF3": "f", # Point<float, 3> - use float for components
66+
}
67+
68+
import os
69+
# Platform-specific adjustments for Windows
70+
if os.name == 'nt':
71+
_format_map['UL'] = 'I' # unsigned int on Windows
72+
_format_map['SL'] = 'i' # signed int on Windows
73+
74+
try:
75+
return _format_map[itk_Image_type]
76+
except KeyError as e:
77+
raise ValueError(f"Unknown ITK image type: {itk_Image_type}") from e
78+
3279
def _get_numpy_pixelid(itk_Image_type) -> np.dtype:
3380
"""Returns a ITK PixelID given a numpy array."""
3481

Modules/Core/Common/wrapping/test/itkImageTest.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
# ==========================================================================
1818
import itk
1919
import numpy as np
20+
import sys
2021

2122
Dimension = 2
2223
PixelType = itk.UC
@@ -30,12 +31,19 @@
3031
image.Allocate()
3132
image.FillBuffer(4)
3233

33-
array = image.__array__()
34+
if sys.version_info >= (3, 12):
35+
array = np.array(image)
36+
else:
37+
array = np.array(image.__buffer__())
38+
3439
assert array[0, 0] == 4
3540
assert array[0, 1] == 4
3641
assert isinstance(array, np.ndarray)
3742

38-
array = np.asarray(image)
43+
if sys.version_info >= (3, 12):
44+
array = np.asarray(image)
45+
else:
46+
array = np.asarray(image.__buffer__())
3947
assert array[0, 0] == 4
4048
assert array[0, 1] == 4
4149
assert isinstance(array, np.ndarray)

Modules/Core/Transform/wrapping/test/itkTransformSerializationTest.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,14 +174,28 @@
174174
convert_filter.SetReferenceImage(fixed_image)
175175
convert_filter.Update()
176176
field1 = convert_filter.GetOutput()
177+
field1.Update()
177178
field1 = np.array(field1)
178179

180+
print(field1)
181+
179182
convert_filter = itk.TransformToDisplacementFieldFilter.IVF22D.New()
180183
convert_filter.SetTransform(serialize_deserialize)
181184
convert_filter.UseReferenceImageOn()
182185
convert_filter.SetReferenceImage(fixed_image)
183186
convert_filter.Update()
184187
field2 = convert_filter.GetOutput()
188+
field2.Update()
185189
field2 = np.array(field2)
186190

187-
assert np.array_equal(np.array(field1), np.array(field2))
191+
192+
if not np.array_equal(field1, field2):
193+
print(f"Field1 shape: {field1.shape}, dtype: {field1.dtype}")
194+
print(f"Field2 shape: {field2.shape}, dtype: {field2.dtype}")
195+
print(f"Field1 min: {field1.min()}, max: {field1.max()}, mean: {field1.mean()}")
196+
print(f"Field2 min: {field2.min()}, max: {field2.max()}, mean: {field2.mean()}")
197+
print(f"Max absolute difference: {np.abs(field1 - field2).max()}")
198+
print(f"Number of different elements: {np.sum(field1 != field2)}")
199+
if field1.shape == field2.shape:
200+
print(f"Are they close (atol=1e-6)? {np.allclose(field1, field2, atol=1e-6)}")
201+
assert np.array_equal(field1, field2)

Wrapping/Generators/Python/PyBase/pyBase.i

Lines changed: 60 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -682,13 +682,67 @@ str = str
682682
683683
%define DECL_PYTHON_IMAGE_CLASS(swig_name)
684684
%extend swig_name {
685-
%pythoncode {
686-
def __array__(self, dtype=None):
685+
%pythoncode %{
686+
def __buffer__(self, flags = 0, / ) -> memoryview:
687687
import itk
688-
import numpy as np
689-
array = itk.array_from_image(self)
690-
return np.asarray(array, dtype=dtype)
691-
}
688+
from itk.itkPyBufferPython import _get_formatstring
689+
690+
itksize = self.GetBufferedRegion().GetSize()
691+
692+
shape = list(itksize)
693+
if self.GetNumberOfComponentsPerPixel() > 1 or isinstance(self, itk.VectorImage):
694+
shape = shape.append(self.GetNumberOfComponentsPerPixel())
695+
696+
shape.reverse()
697+
698+
container_template = itk.template(self)
699+
if container_template is None:
700+
raise BufferError("Cannot determine template parameters for ImportImageContainer")
701+
702+
# Extract pixel type (second template parameter)
703+
pixel_code = container_template[1][0].short_name
704+
format = _get_formatstring(pixel_code)
705+
706+
memview = memoryview(self.GetPixelContainer())
707+
708+
return memview.cast(format, shape=shape)
709+
710+
%}
711+
}
712+
%enddef
713+
714+
%define DECL_PYTHON_IMPORTIMAGECONTAINER_CLASS(swig_name)
715+
%extend swig_name {
716+
%pythoncode %{
717+
def __buffer__(self, flags = 0, / ) -> memoryview:
718+
"""Return a buffer interface for the container.
719+
720+
This allows ImportImageContainer to be used with Python's buffer protocol,
721+
enabling direct memory access from NumPy and other buffer-aware libraries.
722+
"""
723+
import itk
724+
# Get the pixel type from the container template parameters
725+
# The container is templated as ImportImageContainer<IdentifierType, PixelType>
726+
container_template = itk.template(self)
727+
if container_template is None:
728+
raise BufferError("Cannot determine template parameters for ImportImageContainer")
729+
730+
# Extract pixel type (second template parameter)
731+
pixel_type = container_template[1][1]
732+
733+
734+
# Call the PyBuffer method to get the memory view
735+
# We need to determine the appropriate PyBuffer type
736+
try:
737+
# Try to get the PyBuffer class for this pixel type
738+
# PyBuffer is templated over Image types, but we can use a dummy 1D image type
739+
ImageType = itk.Image[pixel_type, 2]
740+
PyBufferType = itk.PyBuffer[ImageType]
741+
742+
return PyBufferType._GetMemoryViewFromImportImageContainer(self)
743+
except (AttributeError, KeyError) as e:
744+
raise BufferError(f"PyBuffer not available for this pixel type: {e}")
745+
%}
692746
}
693747
%enddef
694748

0 commit comments

Comments
 (0)