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..f6965c34b2c 100644 --- a/Wrapping/Generators/Python/PyBase/pyBase.i +++ b/Wrapping/Generators/Python/PyBase/pyBase.i @@ -679,13 +679,58 @@ 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] + 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 + + 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