diff --git a/Modules/Bridge/NumPy/include/itkPyBuffer.h b/Modules/Bridge/NumPy/include/itkPyBuffer.h index 9f294ed4a55..54c1ef3978c 100644 --- a/Modules/Bridge/NumPy/include/itkPyBuffer.h +++ b/Modules/Bridge/NumPy/include/itkPyBuffer.h @@ -72,6 +72,12 @@ class PyBuffer static PyObject * _GetArrayViewFromImage(ImageType * image); + /** + * Get an 1-D byte MemoryView of the container's buffer + */ + static PyObject * + _GetMemoryViewFromImportImageContainer(typename ImageType::PixelContainer * container); + /** * Get an ITK image from a contiguous Python array. Internal helper function for the implementation of * `itkPyBuffer.GetImageViewFromArray`. diff --git a/Modules/Bridge/NumPy/include/itkPyBuffer.hxx b/Modules/Bridge/NumPy/include/itkPyBuffer.hxx index 064c3a9d6b5..10cbd9bb119 100644 --- a/Modules/Bridge/NumPy/include/itkPyBuffer.hxx +++ b/Modules/Bridge/NumPy/include/itkPyBuffer.hxx @@ -50,6 +50,38 @@ PyBuffer::_GetArrayViewFromImage(ImageType * image) return PyMemoryView_FromBuffer(&pyBuffer); } +template +PyObject * +PyBuffer::_GetMemoryViewFromImportImageContainer(typename ImageType::PixelContainer * container) +{ + using ContainerType = typename ImageType::PixelContainer; + Py_buffer pyBuffer{}; + + if (!container) + { + throw std::runtime_error("Input container is null"); + } + + void * const buffer = container->GetBufferPointer(); + + if (!buffer) + { + throw std::runtime_error("Container buffer pointer is null"); + } + + // If the container does not own the buffer then issue a warning + if (!container->GetContainerManageMemory()) + { + PyErr_WarnEx(PyExc_RuntimeWarning, "The ImportImageContainer does not own the exported buffer.", 1); + } + + const SizeValueType size = container->Size(); + const auto len = static_cast(size * sizeof(typename ContainerType::Element)); + + PyBuffer_FillInfo(&pyBuffer, nullptr, buffer, len, 0, PyBUF_CONTIG); + return PyMemoryView_FromBuffer(&pyBuffer); +} + template auto PyBuffer::_get_image_view_from_contiguous_array(PyObject * arr, PyObject * shape, PyObject * numOfComponent) diff --git a/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init b/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init index 8f502c6fd66..d05c074d08c 100644 --- a/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init +++ b/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init @@ -29,6 +29,53 @@ else: loads = dask_deserialize.dispatch(np.ndarray) return NDArrayITKBase(loads(header, frames)) +def _get_formatstring(itk_Image_type) -> str: + """Returns the struct format string for a given ITK image type. + + Format characters from Python's struct module: + - 'b': signed char (int8) + - 'B': unsigned char (uint8) + - 'h': short (int16) + - 'H': unsigned short (uint16) + - 'i': int (int32) + - 'I': unsigned int (uint32) + - 'l': long (platform dependent) + - 'L': unsigned long (platform dependent) + - 'q': long long (int64) + - 'Q': unsigned long long (uint64) + - 'f': float (float32) + - 'd': double (float64) + """ + + # Mapping from ITK pixel type codes to struct format strings + _format_map = { + "UC": "B", # unsigned char + "US": "H", # unsigned short + "UI": "I", # unsigned int + "UL": "L", # unsigned long + "ULL": "Q", # unsigned long long + "SC": "b", # signed char + "SS": "h", # signed short + "SI": "i", # signed int + "SL": "l", # signed long + "SLL": "q", # signed long long + "F": "f", # float + "D": "d", # double + "PF2": "f", # Point - use float for components + "PF3": "f", # Point - use float for components + } + + import os + # Platform-specific adjustments for Windows + if os.name == 'nt': + _format_map['UL'] = 'I' # unsigned int on Windows + _format_map['SL'] = 'i' # signed int on Windows + + try: + return _format_map[itk_Image_type] + except KeyError as e: + raise ValueError(f"Unknown ITK image type: {itk_Image_type}") 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..24bbb2c449e 100644 --- a/Modules/Core/Common/wrapping/test/itkImageTest.py +++ b/Modules/Core/Common/wrapping/test/itkImageTest.py @@ -17,6 +17,7 @@ # ========================================================================== import itk import numpy as np +import sys Dimension = 2 PixelType = itk.UC @@ -30,12 +31,19 @@ image.Allocate() image.FillBuffer(4) -array = image.__array__() +if sys.version_info >= (3, 12): + array = np.array(image) +else: + array = np.array(image.__buffer__()) + assert array[0, 0] == 4 assert array[0, 1] == 4 assert isinstance(array, np.ndarray) -array = np.asarray(image) +if sys.version_info >= (3, 12): + array = np.asarray(image) +else: + array = np.asarray(image.__buffer__()) assert array[0, 0] == 4 assert array[0, 1] == 4 assert isinstance(array, np.ndarray) diff --git a/Modules/Core/Transform/wrapping/test/itkTransformSerializationTest.py b/Modules/Core/Transform/wrapping/test/itkTransformSerializationTest.py index 552ddb75261..1fa032dd01b 100644 --- a/Modules/Core/Transform/wrapping/test/itkTransformSerializationTest.py +++ b/Modules/Core/Transform/wrapping/test/itkTransformSerializationTest.py @@ -174,14 +174,28 @@ convert_filter.SetReferenceImage(fixed_image) convert_filter.Update() field1 = convert_filter.GetOutput() +field1.Update() field1 = np.array(field1) +print(field1) + convert_filter = itk.TransformToDisplacementFieldFilter.IVF22D.New() convert_filter.SetTransform(serialize_deserialize) convert_filter.UseReferenceImageOn() convert_filter.SetReferenceImage(fixed_image) convert_filter.Update() field2 = convert_filter.GetOutput() +field2.Update() field2 = np.array(field2) -assert np.array_equal(np.array(field1), np.array(field2)) + +if not np.array_equal(field1, field2): + print(f"Field1 shape: {field1.shape}, dtype: {field1.dtype}") + print(f"Field2 shape: {field2.shape}, dtype: {field2.dtype}") + print(f"Field1 min: {field1.min()}, max: {field1.max()}, mean: {field1.mean()}") + print(f"Field2 min: {field2.min()}, max: {field2.max()}, mean: {field2.mean()}") + print(f"Max absolute difference: {np.abs(field1 - field2).max()}") + print(f"Number of different elements: {np.sum(field1 != field2)}") + if field1.shape == field2.shape: + print(f"Are they close (atol=1e-6)? {np.allclose(field1, field2, atol=1e-6)}") +assert np.array_equal(field1, field2) diff --git a/Wrapping/Generators/Python/PyBase/pyBase.i b/Wrapping/Generators/Python/PyBase/pyBase.i index 8d695b87aac..04da0d8b2b8 100644 --- a/Wrapping/Generators/Python/PyBase/pyBase.i +++ b/Wrapping/Generators/Python/PyBase/pyBase.i @@ -682,13 +682,67 @@ str = str %define DECL_PYTHON_IMAGE_CLASS(swig_name) %extend swig_name { - %pythoncode { - def __array__(self, dtype=None): + %pythoncode %{ + def __buffer__(self, flags = 0, / ) -> memoryview: import itk - import numpy as np - array = itk.array_from_image(self) - return np.asarray(array, dtype=dtype) - } + from itk.itkPyBufferPython import _get_formatstring + + itksize = self.GetBufferedRegion().GetSize() + + shape = list(itksize) + if self.GetNumberOfComponentsPerPixel() > 1 or isinstance(self, itk.VectorImage): + shape = shape.append(self.GetNumberOfComponentsPerPixel()) + + shape.reverse() + + container_template = itk.template(self) + if container_template is None: + raise BufferError("Cannot determine template parameters for ImportImageContainer") + + # Extract pixel type (second template parameter) + pixel_code = container_template[1][0].short_name + format = _get_formatstring(pixel_code) + + memview = memoryview(self.GetPixelContainer()) + + return memview.cast(format, shape=shape) + + %} + } +%enddef + +%define DECL_PYTHON_IMPORTIMAGECONTAINER_CLASS(swig_name) + %extend swig_name { + %pythoncode %{ + def __buffer__(self, flags = 0, / ) -> memoryview: + """Return a buffer interface for the container. + + This allows ImportImageContainer to be used with Python's buffer protocol, + enabling direct memory access from NumPy and other buffer-aware libraries. + """ + import itk + # Get the pixel type from the container template parameters + # The container is templated as ImportImageContainer + container_template = itk.template(self) + if container_template is None: + raise BufferError("Cannot determine template parameters for ImportImageContainer") + + # Extract pixel type (second template parameter) + pixel_type = container_template[1][1] + + + # Call the PyBuffer method to get the memory view + # We need to determine the appropriate PyBuffer type + try: + # Try to get the PyBuffer class for this pixel type + # PyBuffer is templated over Image types, but we can use a dummy 1D image type + ImageType = itk.Image[pixel_type, 2] + PyBufferType = itk.PyBuffer[ImageType] + + return PyBufferType._GetMemoryViewFromImportImageContainer(self) + except (AttributeError, KeyError) as e: + raise BufferError(f"PyBuffer not available for this pixel type: {e}") + %} } %enddef