From 476e076ba5cc197e81c80c25747095ea0cc1cbb9 Mon Sep 17 00:00:00 2001 From: Hans Johnson Date: Tue, 7 Apr 2026 01:59:30 +0000 Subject: [PATCH 1/2] ENH: Add PEP 688 buffer protocol support for itk.Image MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a ``__buffer__`` method to all wrapped ``itk.Image`` types, enabling zero-copy export of image data as a Python ``memoryview`` with correct shape and element format. On Python 3.12+ (PEP 688), ``memoryview(image)`` and ``numpy.asarray(image)`` automatically call ``__buffer__``. On Python 3.10-3.11 the method can be called explicitly via ``image.__buffer__()``. Key implementation details: - The memoryview is obtained from the existing C++ ``PyBuffer::_GetArrayViewFromImage`` method (zero-copy). - Shape follows NumPy C-order convention (e.g. [z, y, x] for 3D). - Multi-component pixels (RGB, Vector, etc.) get an extra trailing dimension for the components (channels-last convention). - Component type is resolved through ITK template introspection, correctly handling composite pixel types. - A ``_get_buffer_formatstring`` helper maps ITK pixel codes to Python struct format characters, with Windows platform handling. - The existing ``__array__`` method is unchanged — it continues to use ``array_from_image`` for full compatibility. Tests cover scalar images (UC, F, SS), 2D and 3D, shared-memory verification, and automatic PEP 688 dispatch on Python 3.12+. --- Modules/Bridge/NumPy/wrapping/PyBuffer.i.init | 48 +++++++++++++++++++ .../Core/Common/wrapping/test/itkImageTest.py | 43 +++++++++++++++++ Wrapping/Generators/Python/PyBase/pyBase.i | 48 ++++++++++++++++++- 3 files changed, 137 insertions(+), 2 deletions(-) diff --git a/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init b/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init index 8f502c6fd66..1dab77c0242 100644 --- a/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init +++ b/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init @@ -29,6 +29,54 @@ else: loads = dask_deserialize.dispatch(np.ndarray) return NDArrayITKBase(loads(header, frames)) +def _get_buffer_formatstring(itk_pixel_code: str) -> str: + """Return the struct format character for an ITK pixel type code. + + Used by the PEP 688 ``__buffer__`` protocol implementation on + ``itk.Image`` to describe the element type of the exported + memoryview. Format characters follow Python's ``struct`` module + specification. + + Parameters + ---------- + itk_pixel_code : str + Short name of the ITK component type, e.g. ``"UC"``, ``"F"``. + + Returns + ------- + str + Single-character ``struct`` format string. + """ + + _format_map = { + "UC": "B", # unsigned char -> uint8 + "US": "H", # unsigned short -> uint16 + "UI": "I", # unsigned int -> uint32 + "UL": "L", # unsigned long -> platform + "ULL": "Q", # unsigned long long -> uint64 + "SC": "b", # signed char -> int8 + "SS": "h", # short -> int16 + "SI": "i", # int -> int32 + "SL": "l", # long -> platform + "SLL": "q", # long long -> int64 + "F": "f", # float -> float32 + "D": "d", # double -> float64 + } + + import os + if os.name == 'nt': + # On Windows, C ``long`` is 32-bit + _format_map['UL'] = 'I' + _format_map['SL'] = 'i' + + try: + return _format_map[itk_pixel_code] + except KeyError as e: + raise ValueError( + f"No struct format for ITK pixel code '{itk_pixel_code}'" + ) from e + + def _get_numpy_pixelid(itk_Image_type) -> np.dtype: """Returns a ITK PixelID given a numpy array.""" diff --git a/Modules/Core/Common/wrapping/test/itkImageTest.py b/Modules/Core/Common/wrapping/test/itkImageTest.py index 140806f119a..559d59d724f 100644 --- a/Modules/Core/Common/wrapping/test/itkImageTest.py +++ b/Modules/Core/Common/wrapping/test/itkImageTest.py @@ -15,6 +15,7 @@ # limitations under the License. # # ========================================================================== +import sys import itk import numpy as np @@ -39,3 +40,45 @@ assert array[0, 0] == 4 assert array[0, 1] == 4 assert isinstance(array, np.ndarray) + +# --- PEP 688 buffer protocol tests --- + +# Test __buffer__ on scalar image +mv = image.__buffer__() +assert isinstance(mv, memoryview) +assert mv.format == "B" # unsigned char +assert mv.shape == (10, 10) +assert mv[0, 0] == 4 +assert mv.readonly is False + +# Test that memoryview shares memory (zero-copy) +image.SetPixel([3, 5], 42) +assert mv[5, 3] == 42 + +# Test float image +float_image = itk.Image[itk.F, 2].New() +float_image.SetRegions([8, 6]) +float_image.Allocate() +float_image.FillBuffer(3.14) +fmv = float_image.__buffer__() +assert fmv.format == "f" +assert fmv.shape == (6, 8) +assert abs(fmv[0, 0] - 3.14) < 1e-5 + +# Test 3D image +image_3d = itk.Image[itk.SS, 3].New() +image_3d.SetRegions([4, 5, 6]) +image_3d.Allocate() +image_3d.FillBuffer(-7) +mv3d = image_3d.__buffer__() +assert mv3d.format == "h" # signed short +assert mv3d.shape == (6, 5, 4) +assert mv3d[0, 0, 0] == -7 + +# On Python 3.12+, memoryview() should call __buffer__ automatically +if sys.version_info >= (3, 12): + auto_mv = memoryview(image) + assert auto_mv.shape == (10, 10) + assert auto_mv[0, 0] == 4 + +print("All buffer protocol tests passed.") diff --git a/Wrapping/Generators/Python/PyBase/pyBase.i b/Wrapping/Generators/Python/PyBase/pyBase.i index 1f825de6434..91ea059dc81 100644 --- a/Wrapping/Generators/Python/PyBase/pyBase.i +++ b/Wrapping/Generators/Python/PyBase/pyBase.i @@ -679,13 +679,57 @@ str = str %define DECL_PYTHON_IMAGE_CLASS(swig_name) %extend swig_name { - %pythoncode { + %pythoncode %{ + def __buffer__(self, flags=0, /): + """PEP 688 buffer protocol — export image data as a memoryview. + + On Python 3.12+ this is called automatically by + ``memoryview(image)`` and ``numpy.asarray(image)``. + On Python 3.10–3.11 it can be called explicitly. + + The returned memoryview shares memory with the image + (zero-copy). The caller must keep the image alive for + as long as the memoryview is in use. + """ + import itk + from itk.itkPyBufferPython import _get_buffer_formatstring + + # Get 1-D raw memoryview from the C++ buffer + ImageType = type(self) + PyBufferType = itk.PyBuffer[ImageType] + raw_memview = PyBufferType._GetArrayViewFromImage(self) + + # Build shape in C-order (NumPy convention: [z, y, x, ...]) + itksize = self.GetBufferedRegion().GetSize() + shape = [int(itksize[d]) for d in range(len(itksize))] + + n_components = self.GetNumberOfComponentsPerPixel() + if n_components > 1 or isinstance(self, itk.VectorImage): + shape.insert(0, n_components) + + shape.reverse() + + # Determine the struct format character for the component type + tpl = itk.template(self) + pixel_type = tpl[1][0] + pixel_tpl = itk.template(pixel_type) + if pixel_tpl and len(pixel_tpl[1]) > 0: + # Composite pixel (RGB, Vector, etc.) — use component type + component_code = pixel_tpl[1][0].short_name + else: + # Scalar pixel + component_code = pixel_type.short_name + + fmt = _get_buffer_formatstring(component_code) + + return raw_memview.cast(fmt, shape=shape) + def __array__(self, dtype=None): import itk import numpy as np array = itk.array_from_image(self) return np.asarray(array, dtype=dtype) - } + %} } %enddef From 4f0c6764d97835a70aba5503e8a0a04f47bd5ee6 Mon Sep 17 00:00:00 2001 From: Hans Johnson Date: Tue, 7 Apr 2026 03:05:45 +0000 Subject: [PATCH 2/2] BUG: Fix __buffer__ pixel type detection for scalar images Handle itkCType (scalar pixel types like UC, F, SS) directly via short_name instead of calling itk.template() which fails on non-template types. Only use itk.template() for composite pixel types (RGB, Vector, etc.). Co-Authored-By: Claude Opus 4.6 (1M context) --- Wrapping/Generators/Python/PyBase/pyBase.i | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/Wrapping/Generators/Python/PyBase/pyBase.i b/Wrapping/Generators/Python/PyBase/pyBase.i index 91ea059dc81..f6965c34b2c 100644 --- a/Wrapping/Generators/Python/PyBase/pyBase.i +++ b/Wrapping/Generators/Python/PyBase/pyBase.i @@ -712,13 +712,14 @@ str = str # Determine the struct format character for the component type tpl = itk.template(self) pixel_type = tpl[1][0] - pixel_tpl = itk.template(pixel_type) - if pixel_tpl and len(pixel_tpl[1]) > 0: + from itk.support.types import itkCType + if isinstance(pixel_type, itkCType): + # Scalar pixel (UC, F, SS, etc.) + component_code = pixel_type.short_name + else: # Composite pixel (RGB, Vector, etc.) — use component type + pixel_tpl = itk.template(pixel_type) component_code = pixel_tpl[1][0].short_name - else: - # Scalar pixel - component_code = pixel_type.short_name fmt = _get_buffer_formatstring(component_code)