@@ -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