Skip to content

Commit edecc6c

Browse files
hjmjohnsonhjmjohnson
committed
ENH: Add PEP 688 buffer protocol and fix np.asarray() lifetime safety
Add zero-copy data export to all wrapped itk.Image types via two protocols, ensuring the exported array remains valid even after the source image is deleted: image = itk.imread("brain.nii.gz") arr = np.asarray(image) del image print(arr[1,1,1]) # safe -- no crash Protocol dispatch by Python version: 3.12+: np.asarray -> __buffer__ (PEP 688, zero-copy, memoryview pins self via NDArrayITKBase intermediary) 3.10-11: np.asarray -> __array__ -> array_view_from_image (zero-copy, NDArrayITKBase.itk_base holds reference to image) Changes to pyBase.i: - Add __buffer__() implementing PEP 688 buffer export with shaped memoryview. Uses NDArrayITKBase as intermediary to hold a Python reference to the image, preventing GC while any derived memoryview/array exists. - Simplify __array__() to always return zero-copy view via array_view_from_image(). Supports NumPy 2.0 copy= parameter. copy=True returns a plain ndarray (not NDArrayITKBase) so the image can be GCd immediately. - Remove __array_interface__ (returned raw pointer with no reference holder -- use-after-free on del image, confirmed by test). - Remove SIMULATE_PEP688 / SIMULATE_PEP688_DEBUG (confusing, contradictory behaviors between __buffer__ and __array__ paths). Changes to PyBuffer.i.init: - Add _get_buffer_formatstring() with module-level _BUFFER_FORMAT_MAP - Add _get_numpy_pixelid() with module-level _NUMPY_PIXELID_MAP - Remove LD (long double) mapping (silent corruption) Test suites added (121 assertions): itkImageTest.py (29): __buffer__, memoryview, np.asarray, __array__ itkImageInteropTest.py (60): NumPy, PyTorch, Dask clinical sizes itkImageLifetimeTest.py (32): del image on every export path Supersedes InsightSoftwareConsortium#6020, InsightSoftwareConsortium#6018, InsightSoftwareConsortium#5673, InsightSoftwareConsortium#5665. Co-Authored-By: Hans J. Johnson <hans-johnson@uiowa.edu>
1 parent 35158b6 commit edecc6c

6 files changed

Lines changed: 1181 additions & 29 deletions

File tree

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

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

32+
# Module-level constant — built once at import time.
33+
# Composite pixel types (RGB, Vector, etc.) are decomposed to their
34+
# scalar component type before reaching _get_buffer_formatstring(),
35+
# so only scalar codes appear here.
36+
# Platform aliases (ST/IT/OT) resolve to UL/SL at the wrapping level.
37+
import os as _os
38+
_BUFFER_FORMAT_MAP = {
39+
# --- integer types ---
40+
"B": "?", # bool
41+
"UC": "B", # unsigned char -> uint8
42+
"US": "H", # unsigned short -> uint16
43+
"UI": "I", # unsigned int -> uint32
44+
"UL": "L", # unsigned long -> platform (8 bytes Linux/macOS, 4 bytes Windows)
45+
"ULL": "Q", # unsigned long long -> uint64
46+
"SC": "b", # signed char -> int8
47+
"SS": "h", # signed short -> int16
48+
"SI": "i", # signed int -> int32
49+
"SL": "l", # signed long -> platform
50+
"SLL": "q", # signed long long -> int64
51+
# --- floating point types ---
52+
"F": "f", # float -> float32
53+
"D": "d", # double -> float64
54+
# "LD" (long double) intentionally omitted: sizeof(long double)
55+
# is 16 bytes on Linux/macOS x86-64 but struct "d" is 8 bytes.
56+
# Casting would silently corrupt the buffer.
57+
}
58+
59+
def _get_buffer_formatstring(itk_pixel_code: str) -> str:
60+
"""Return the struct format character for an ITK pixel type code.
61+
62+
Used by the PEP 688 ``__buffer__`` protocol implementation on
63+
``itk.Image`` to describe the element type of the exported
64+
memoryview. Format characters follow Python's ``struct`` module
65+
specification.
66+
67+
Parameters
68+
----------
69+
itk_pixel_code : str
70+
Short name of the ITK component type, e.g. ``"UC"``, ``"F"``.
71+
72+
Returns
73+
-------
74+
str
75+
Single-character ``struct`` format string.
76+
77+
Raises
78+
------
79+
KeyError
80+
If ``itk_pixel_code`` is not a supported scalar type.
81+
"""
82+
try:
83+
return _BUFFER_FORMAT_MAP[itk_pixel_code]
84+
except KeyError:
85+
raise KeyError(
86+
f"Unsupported ITK pixel type code {itk_pixel_code!r} for "
87+
f"buffer export. Supported codes: "
88+
f"{sorted(_BUFFER_FORMAT_MAP)}"
89+
) from None
90+
if _os.name == 'nt':
91+
# On Windows, C ``long`` is 32-bit
92+
_BUFFER_FORMAT_MAP['UL'] = 'I'
93+
_BUFFER_FORMAT_MAP['SL'] = 'i'
94+
95+
3296
def _get_numpy_pixelid(itk_Image_type) -> np.dtype:
3397
"""Returns a ITK PixelID given a numpy array."""
3498

35-
# This is a Mapping from numpy array types to itk pixel types.
36-
_np_itk = {"UC":np.dtype(np.uint8),
37-
"US":np.dtype(np.uint16),
38-
"UI":np.dtype(np.uint32),
39-
"UL":np.dtype(np.uint64),
40-
"ULL":np.dtype(np.uint64),
41-
"SC":np.dtype(np.int8),
42-
"SS":np.dtype(np.int16),
43-
"SI":np.dtype(np.int32),
44-
"SL":np.dtype(np.int64),
45-
"SLL":np.dtype(np.int64),
46-
"F":np.dtype(np.float32),
47-
"D":np.dtype(np.float64),
48-
"PF2":np.dtype(np.float32),
49-
"PF3":np.dtype(np.float32),
50-
}
51-
import os
52-
if os.name == 'nt':
53-
_np_itk['UL'] = np.dtype(np.uint32)
54-
_np_itk['SL'] = np.dtype(np.int32)
55-
try:
56-
return _np_itk[itk_Image_type]
57-
except KeyError as e:
58-
raise e
99+
return _NUMPY_PIXELID_MAP[itk_Image_type]
100+
101+
102+
# Module-level constant — built once at import time.
103+
_NUMPY_PIXELID_MAP = {
104+
"B": np.dtype(np.bool_),
105+
"UC": np.dtype(np.uint8),
106+
"US": np.dtype(np.uint16),
107+
"UI": np.dtype(np.uint32),
108+
"UL": np.dtype(np.uint64),
109+
"ULL": np.dtype(np.uint64),
110+
"SC": np.dtype(np.int8),
111+
"SS": np.dtype(np.int16),
112+
"SI": np.dtype(np.int32),
113+
"SL": np.dtype(np.int64),
114+
"SLL": np.dtype(np.int64),
115+
"F": np.dtype(np.float32),
116+
"D": np.dtype(np.float64),
117+
"PF2": np.dtype(np.float32),
118+
"PF3": np.dtype(np.float32),
119+
}
120+
if _os.name == 'nt':
121+
_NUMPY_PIXELID_MAP['UL'] = np.dtype(np.uint32)
122+
_NUMPY_PIXELID_MAP['SL'] = np.dtype(np.int32)
59123
%}

‎Modules/Core/Common/wrapping/test/CMakeLists.txt‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,4 +56,14 @@ if(ITK_WRAP_PYTHON)
5656
COMMAND
5757
${CMAKE_CURRENT_SOURCE_DIR}/itkSpatialOrientationAdapterTest.py
5858
)
59+
itk_python_add_test(
60+
NAME itkImageInteropPythonTest
61+
COMMAND
62+
${CMAKE_CURRENT_SOURCE_DIR}/itkImageInteropTest.py
63+
)
64+
itk_python_add_test(
65+
NAME itkImageLifetimePythonTest
66+
COMMAND
67+
${CMAKE_CURRENT_SOURCE_DIR}/itkImageLifetimeTest.py
68+
)
5969
endif()

0 commit comments

Comments
 (0)