diff --git a/Modules/Core/Common/wrapping/test/CMakeLists.txt b/Modules/Core/Common/wrapping/test/CMakeLists.txt index 76ae8305118..2f948e1d43c 100644 --- a/Modules/Core/Common/wrapping/test/CMakeLists.txt +++ b/Modules/Core/Common/wrapping/test/CMakeLists.txt @@ -56,4 +56,14 @@ if(ITK_WRAP_PYTHON) COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/itkSpatialOrientationAdapterTest.py ) + itk_python_add_test( + NAME itkImageInteropPythonTest + COMMAND + ${CMAKE_CURRENT_SOURCE_DIR}/itkImageInteropTest.py + ) + itk_python_add_test( + NAME itkImageLifetimePythonTest + COMMAND + ${CMAKE_CURRENT_SOURCE_DIR}/itkImageLifetimeTest.py + ) endif() diff --git a/Modules/Core/Common/wrapping/test/itkImageInteropTest.py b/Modules/Core/Common/wrapping/test/itkImageInteropTest.py new file mode 100755 index 00000000000..8e341621408 --- /dev/null +++ b/Modules/Core/Common/wrapping/test/itkImageInteropTest.py @@ -0,0 +1,454 @@ +#!/usr/bin/env python +# ========================================================================== +# +# Copyright NumFOCUS +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0.txt +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# ========================================================================== +"""Comprehensive interop tests for itk.Image with numpy, torch, and dask. + +Tests use image sizes representative of clinical imaging modalities: + - CT: 512x512x128 signed short (int16) + - MRI: 256x256x64 float32 + - DWI: 128x128x40 float32 vector (30 gradient directions) + +Required: numpy +Optional: torch, dask (tests skipped if unavailable) +""" +import sys +import itk +import numpy as np + +# --- Optional dependency detection --- +_HAVE_TORCH = False +try: + import torch + + _HAVE_TORCH = True +except ImportError: + pass + +_HAVE_DASK = False +try: + import dask + import dask.array as da + + _HAVE_DASK = True +except ImportError: + pass + +passed = 0 +skipped = 0 +failed = 0 + + +def check(name, condition, detail=""): + global passed, failed + if condition: + passed += 1 + else: + failed += 1 + msg = f" FAIL: {name}" + if detail: + msg += f" ({detail})" + print(msg) + + +def skip(name, reason): + global skipped + skipped += 1 + print(f" SKIP: {name} ({reason})") + + +# --- Feature detection --- +_probe = itk.Image[itk.SS, 3].New() +_probe.SetRegions([2, 2, 2]) +_probe.Allocate() +HAS_BUFFER = hasattr(_probe, "__buffer__") + +# Detect whether __array__ accepts the copy keyword (NumPy 2.x protocol). +HAS_ARRAY_COPY_PARAM = False +try: + _probe.__array__(copy=None) + HAS_ARRAY_COPY_PARAM = True +except TypeError: + pass + +# Detect whether np.asarray produces a zero-copy view. +_probe.FillBuffer(0) +_a = np.asarray(_probe) +_probe.SetPixel([0, 0, 0], 999) +HAS_ZEROCOPY = _a[0, 0, 0] == 999 +del _a, _probe + + +# =================================================================== +# Helper: create images matching clinical modalities +# =================================================================== + + +def make_ct_image(): + """CT: 512x512x128, signed short (int16), ~67 MB.""" + img = itk.Image[itk.SS, 3].New() + img.SetRegions([512, 512, 128]) + img.SetSpacing([0.5, 0.5, 1.0]) + img.Allocate() + img.FillBuffer(-1024) # Hounsfield air + return img + + +def make_mri_image(): + """MRI T1: 256x256x64, float32, ~16 MB.""" + img = itk.Image[itk.F, 3].New() + img.SetRegions([256, 256, 64]) + img.SetSpacing([1.0, 1.0, 2.0]) + img.Allocate() + img.FillBuffer(500.0) # Typical T1 signal + return img + + +def make_dwi_image(): + """DWI: 128x128x40, float32, single volume, ~2.5 MB.""" + img = itk.Image[itk.F, 3].New() + img.SetRegions([128, 128, 40]) + img.SetSpacing([2.0, 2.0, 2.5]) + img.Allocate() + img.FillBuffer(800.0) # B0 signal + return img + + +def make_rgb_image(): + """RGB color: 256x256, unsigned char, ~192 KB.""" + img = itk.Image[itk.RGBPixel[itk.UC], 2].New() + img.SetRegions([256, 256]) + img.Allocate() + return img + + +def make_vector_image(): + """Vector image: 64x64x32, 6 components (DTI tensor), float32.""" + img = itk.VectorImage[itk.F, 3].New() + img.SetRegions([64, 64, 32]) + img.SetNumberOfComponentsPerPixel(6) + img.Allocate(True) # zero-initialize + return img + + +# =================================================================== +# Section 1: NumPy interop (REQUIRED) +# =================================================================== +print("=" * 60) +print("Section 1: NumPy interop (required)") +print("=" * 60) + +# -- 1.1 np.asarray basic properties -- +ct = make_ct_image() +ct_arr = np.asarray(ct) +check("CT np.asarray shape", ct_arr.shape == (128, 512, 512)) +check("CT np.asarray dtype", ct_arr.dtype == np.int16) +check("CT np.asarray value", ct_arr[0, 0, 0] == -1024) + +# -- 1.2 np.asarray zero-copy -- +if HAS_ZEROCOPY: + # Verify zero-copy: modify through numpy, read back through ITK + ct_arr[10, 20, 30] = 1500 + check("CT zero-copy np->itk", ct.GetPixel([30, 20, 10]) == 1500) + + # Verify zero-copy: modify through ITK, read back through numpy + ct.SetPixel([5, 6, 7], 42) + check("CT zero-copy itk->np", ct_arr[7, 6, 5] == 42) +else: + skip("CT zero-copy np->itk", "zero-copy np.asarray not available") + skip("CT zero-copy itk->np", "zero-copy np.asarray not available") + +# -- 1.3 MRI image -- +mri = make_mri_image() +mri_arr = np.asarray(mri) +check("MRI np.asarray shape", mri_arr.shape == (64, 256, 256)) +check("MRI np.asarray dtype", mri_arr.dtype == np.float32) +check("MRI np.asarray value", abs(mri_arr[0, 0, 0] - 500.0) < 1e-5) + +# Zero-copy write +if HAS_ZEROCOPY: + mri_arr[32, 128, 128] = 999.0 + check("MRI zero-copy", mri.GetPixel([128, 128, 32]) == 999.0) +else: + skip("MRI zero-copy", "zero-copy np.asarray not available") + +# -- 1.4 DWI single volume -- +dwi = make_dwi_image() +dwi_arr = np.asarray(dwi) +check("DWI np.asarray shape", dwi_arr.shape == (40, 128, 128)) +check("DWI np.asarray dtype", dwi_arr.dtype == np.float32) + +# -- 1.5 Multiple pixel types -- +for ptype, expected_dtype, label in [ + (itk.UC, np.uint8, "UC/uint8"), + (itk.SS, np.int16, "SS/int16"), + (itk.SI, np.int32, "SI/int32"), + (itk.F, np.float32, "F/float32"), + (itk.D, np.float64, "D/float64"), +]: + img = itk.Image[ptype, 3].New() + img.SetRegions([4, 4, 4]) + img.Allocate() + arr = np.asarray(img) + check(f"{label} dtype", arr.dtype == expected_dtype, f"got {arr.dtype}") + check(f"{label} shape", arr.shape == (4, 4, 4)) + +# -- 1.6 __array__ returns view (not copy) -- +if HAS_ZEROCOPY: + mri2 = make_mri_image() + arr1 = mri2.__array__() + check("__array__ returns ndarray", isinstance(arr1, np.ndarray)) + arr1[0, 0, 0] = 12345.0 + check("__array__ zero-copy", mri2.GetPixel([0, 0, 0]) == 12345.0) + del mri2, arr1 +else: + skip("__array__ zero-copy", "zero-copy np.asarray not available") + +# -- 1.6b __array__(copy=...) NumPy 2.x protocol -- +if HAS_ARRAY_COPY_PARAM: + img_copy_test = make_mri_image() + + # copy=None should return a view (zero-copy) + arr_none = img_copy_test.__array__(copy=None) + check("__array__(copy=None) is ndarray", isinstance(arr_none, np.ndarray)) + if HAS_ZEROCOPY: + arr_none[0, 0, 0] = 77777.0 + check( + "__array__(copy=None) zero-copy", + img_copy_test.GetPixel([0, 0, 0]) == 77777.0, + ) + + # copy=True should return a deep copy + img_copy_test.FillBuffer(500.0) + arr_copy = img_copy_test.__array__(copy=True) + arr_copy[0, 0, 0] = 99999.0 + check( + "__array__(copy=True) is deep copy", + img_copy_test.GetPixel([0, 0, 0]) == 500.0, + ) + + # copy=False should return a view (zero-copy, no dtype conversion) + arr_false = img_copy_test.__array__(copy=False) + check("__array__(copy=False) is ndarray", isinstance(arr_false, np.ndarray)) + if HAS_ZEROCOPY: + arr_false[0, 0, 0] = 11111.0 + check( + "__array__(copy=False) zero-copy", + img_copy_test.GetPixel([0, 0, 0]) == 11111.0, + ) + + # copy=False with incompatible dtype should raise ValueError + raised = False + try: + img_copy_test.__array__(dtype=np.float64, copy=False) + except ValueError: + raised = True + check("__array__(dtype=float64, copy=False) raises ValueError", raised) + + del img_copy_test, arr_none, arr_copy, arr_false +else: + skip("__array__(copy=...) tests", "__array__ copy param not available") + +# -- 1.7 __buffer__ protocol -- +if HAS_BUFFER: + ct_mv = ct.__buffer__() + check("CT __buffer__ type", isinstance(ct_mv, memoryview)) + check("CT __buffer__ format", ct_mv.format == "h") # signed short + check("CT __buffer__ shape", ct_mv.shape == (128, 512, 512)) + + if sys.version_info >= (3, 12): + auto_mv = memoryview(ct) + check("PEP 688 auto memoryview", auto_mv.shape == (128, 512, 512)) + else: + skip( + "PEP 688 auto memoryview", + f"requires Python 3.12+, have {sys.version_info[:2]}", + ) +else: + skip("__buffer__ protocol tests", "__buffer__ not available in this ITK version") + +# -- 1.8 Vector/multi-component images -- +# Note: VectorImage uses array_view_from_image path (DECL_PYTHON_IMAGE_CLASS +# is only applied to itk::Image, not itk::VectorImage) +vec = make_vector_image() +vec_arr = itk.array_view_from_image(vec) +check("VectorImage shape", vec_arr.shape == (32, 64, 64, 6)) +check("VectorImage dtype", vec_arr.dtype == np.float32) + +# -- 1.9 np.array with dtype conversion -- +ct_f32 = np.array(ct, dtype=np.float32) +check("np.array dtype conversion", ct_f32.dtype == np.float32) + +# -- 1.10 Consistency: np.asarray vs __buffer__ vs array_view -- +mri3 = make_mri_image() +via_asarray = np.asarray(mri3) +via_view = itk.array_view_from_image(mri3) +if HAS_BUFFER: + via_buf = np.asarray(mri3.__buffer__()) + check( + "Consistency: all paths same data", + np.array_equal(via_asarray, via_view) and np.array_equal(via_asarray, via_buf), + ) +else: + check( + "Consistency: np.asarray vs array_view", + np.array_equal(via_asarray, via_view), + ) + +print(f" NumPy: {passed} passed, {failed} failed") +numpy_passed = passed +numpy_failed = failed + +# =================================================================== +# Section 2: PyTorch interop (optional) +# =================================================================== +print() +print("=" * 60) +if _HAVE_TORCH: + print(f"Section 2: PyTorch interop (torch {torch.__version__})") +else: + print("Section 2: PyTorch interop (SKIPPED — torch not installed)") +print("=" * 60) + +if _HAVE_TORCH: + # -- 2.1 torch.as_tensor via np.asarray -- + ct2 = make_ct_image() + ct_tensor = torch.as_tensor(np.asarray(ct2)) + check("CT torch.as_tensor shape", tuple(ct_tensor.shape) == (128, 512, 512)) + check("CT torch.as_tensor dtype", ct_tensor.dtype == torch.int16) + check("CT torch.as_tensor value", ct_tensor[0, 0, 0].item() == -1024) + + # -- 2.2 torch.from_numpy zero-copy chain -- + if HAS_ZEROCOPY: + mri4 = make_mri_image() + mri_np = np.asarray(mri4) # zero-copy via np.asarray + mri_tensor = torch.from_numpy(mri_np) # zero-copy numpy->torch + check("MRI torch shape", tuple(mri_tensor.shape) == (64, 256, 256)) + check("MRI torch dtype", mri_tensor.dtype == torch.float32) + + # Zero-copy chain: ITK -> numpy -> torch + mri4.SetPixel([100, 100, 30], 777.0) + check( + "MRI zero-copy itk->np->torch", + mri_tensor[30, 100, 100].item() == 777.0, + ) + else: + skip("MRI torch zero-copy chain", "zero-copy np.asarray not available") + + # -- 2.3 DWI volume -- + dwi2 = make_dwi_image() + dwi_tensor = torch.from_numpy(np.asarray(dwi2)) + check("DWI torch shape", tuple(dwi_tensor.shape) == (40, 128, 128)) + + # -- 2.4 Vector image -- + vec2 = make_vector_image() + vec_tensor = torch.from_numpy(itk.array_view_from_image(vec2)) + check("VectorImage torch shape", tuple(vec_tensor.shape) == (32, 64, 64, 6)) + check("VectorImage torch dtype", vec_tensor.dtype == torch.float32) + + # -- 2.5 dtype mapping -- + for ptype, expected_tdtype, label in [ + (itk.UC, torch.uint8, "UC->uint8"), + (itk.SS, torch.int16, "SS->int16"), + (itk.SI, torch.int32, "SI->int32"), + (itk.F, torch.float32, "F->float32"), + (itk.D, torch.float64, "D->float64"), + ]: + img = itk.Image[ptype, 3].New() + img.SetRegions([4, 4, 4]) + img.Allocate() + t = torch.from_numpy(np.asarray(img)) + check(f"torch {label}", t.dtype == expected_tdtype, f"got {t.dtype}") + + torch_passed = passed - numpy_passed + torch_failed = failed - numpy_failed + print(f" PyTorch: {torch_passed} passed, {torch_failed} failed") +else: + skip("PyTorch tests", "torch not installed") + print(" (install torch to enable these tests)") + +# =================================================================== +# Section 3: Dask interop (optional) +# =================================================================== +print() +print("=" * 60) +if _HAVE_DASK: + print(f"Section 3: Dask interop (dask {dask.__version__})") +else: + print("Section 3: Dask interop (SKIPPED — dask not installed)") +print("=" * 60) + +pre_dask_passed = passed +pre_dask_failed = failed + +if _HAVE_DASK: + # -- 3.1 dask.array.from_array via np.asarray -- + ct3 = make_ct_image() + ct_darr = da.from_array(np.asarray(ct3), chunks=(32, 128, 128)) + check("CT dask shape", ct_darr.shape == (128, 512, 512)) + check("CT dask dtype", ct_darr.dtype == np.int16) + + # Compute a chunk and verify value + chunk = ct_darr[0:32, 0:128, 0:128].compute() + check("CT dask compute value", chunk[0, 0, 0] == -1024) + check("CT dask compute shape", chunk.shape == (32, 128, 128)) + + # -- 3.2 MRI with dask -- + mri5 = make_mri_image() + mri_darr = da.from_array(np.asarray(mri5), chunks=(16, 64, 64)) + check("MRI dask shape", mri_darr.shape == (64, 256, 256)) + check("MRI dask dtype", mri_darr.dtype == np.float32) + + # -- 3.3 Dask reduction on clinical-size image -- + mri_mean = mri_darr.mean().compute() + check("MRI dask mean", abs(mri_mean - 500.0) < 1e-3, f"got {mri_mean}") + + # -- 3.4 DWI with dask -- + dwi3 = make_dwi_image() + dwi_darr = da.from_array(np.asarray(dwi3), chunks=(10, 32, 32)) + check("DWI dask shape", dwi_darr.shape == (40, 128, 128)) + + # -- 3.5 Vector image with dask -- + vec3 = make_vector_image() + vec_darr = da.from_array(itk.array_view_from_image(vec3), chunks=(8, 16, 16, 6)) + check("VectorImage dask shape", vec_darr.shape == (32, 64, 64, 6)) + + dask_passed = passed - pre_dask_passed + dask_failed = failed - pre_dask_failed + print(f" Dask: {dask_passed} passed, {dask_failed} failed") +else: + skip("Dask tests", "dask not installed") + print(" (install dask to enable these tests)") + +# =================================================================== +# Summary +# =================================================================== +print() +print("=" * 60) +print(f"Python {sys.version_info[0]}.{sys.version_info[1]}.{sys.version_info[2]}") +print(f"NumPy {np.__version__}") +print( + f"Features: __buffer__={HAS_BUFFER}, copy_param={HAS_ARRAY_COPY_PARAM}, " + f"zero_copy={HAS_ZEROCOPY}" +) +print(f"TOTAL: {passed} passed, {failed} failed, {skipped} skipped") +print("=" * 60) + +if failed > 0: + sys.exit(1) + +print("All interop tests passed.") diff --git a/Modules/Core/Common/wrapping/test/itkImageLifetimeTest.py b/Modules/Core/Common/wrapping/test/itkImageLifetimeTest.py new file mode 100644 index 00000000000..1eadf45773b --- /dev/null +++ b/Modules/Core/Common/wrapping/test/itkImageLifetimeTest.py @@ -0,0 +1,703 @@ +# ========================================================================== +# +# Copyright NumFOCUS +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0.txt +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# ========================================================================== +"""Test object lifetime, memory safety, and leak detection for itk.Image exports. + +Tests are automatically skipped when features are not available in the +current ITK build. This allows cherry-picking into older ITK releases +(e.g., v5.4.5) where PEP 688 and zero-copy np.asarray are not yet +implemented. + +Features detected at runtime: + - __buffer__: PEP 688 buffer export (ITK 6.x) + - __array__(copy=...): NumPy 2.0 copy parameter (ITK 6.x) + - Zero-copy np.asarray: mutation visible through array (ITK 6.x) +""" + +import gc +import os +import sys +import weakref + +import itk +import numpy as np + +passed = 0 +failed = 0 +skipped = 0 + + +def check(name, condition, detail=""): + global passed, failed + if condition: + passed += 1 + else: + failed += 1 + msg = f" FAIL: {name}" + if detail: + msg += f" ({detail})" + print(msg) + + +def skip(name, reason): + global skipped + skipped += 1 + print(f" SKIP: {name} ({reason})") + + +def make_image(pixel_type=itk.F, dimension=3, size=None, fill=42.0): + """Create a test image with known pixel values.""" + if size is None: + size = [16, 16, 8] + ImageType = itk.Image[pixel_type, dimension] + image = ImageType.New() + image.SetRegions(size[:dimension]) + image.Allocate() + image.FillBuffer(fill) + return image + + +def get_rss_bytes(): + """Get current resident set size in bytes (Linux/macOS).""" + try: + with open("/proc/self/statm") as f: + rss_pages = int(f.read().split()[1]) + return rss_pages * os.sysconf("SC_PAGE_SIZE") + except (OSError, ValueError): + pass + try: + import resource + + rusage = resource.getrusage(resource.RUSAGE_SELF) + if sys.platform == "darwin": + return rusage.ru_maxrss + return rusage.ru_maxrss * 1024 + except (ImportError, AttributeError): + return None + + +# --- Feature detection --- +_probe = make_image(size=[2, 2, 2]) +HAS_BUFFER = hasattr(_probe, "__buffer__") +HAS_ARRAY_COPY_PARAM = False +try: + _probe.__array__(copy=None) + HAS_ARRAY_COPY_PARAM = True +except TypeError: + pass +_probe.FillBuffer(1.0) +_a = np.asarray(_probe) +_probe.SetPixel([0, 0, 0], 999.0) +HAS_ZEROCOPY = _a[0, 0, 0] == 999.0 +del _a, _probe +gc.collect() + +_SKIP_BUFFER = "__buffer__ not available in this ITK version" +_SKIP_ZEROCOPY = "zero-copy np.asarray not available in this ITK version" +_SKIP_COPY = "__array__(copy=) not available in this ITK version" + + +def run_lifetime_and_leak_tests(): + """Run all lifetime tests inside a function scope. + + All local variables are explicitly deleted at the end of each + section, and gc.collect() is called to ensure deterministic + cleanup. After this function returns, the caller verifies + that no image memory has leaked. + """ + + # ============================================================== + # Section 1: Deep copy paths (always safe, all versions) + # ============================================================== + print("=" * 60) + print("Section 1: Deep copy paths (always safe)") + print("=" * 60) + + if HAS_ARRAY_COPY_PARAM: + image = make_image() + arr = image.__array__(copy=True) + expected = arr[0, 0, 0] + del image + gc.collect() + check("__array__(copy=True) survives del image", arr[0, 0, 0] == expected) + del arr, expected + else: + skip("__array__(copy=True) survives del image", _SKIP_COPY) + + image = make_image(fill=99.0) + arr = np.array(image, copy=True) + del image + gc.collect() + check("np.array(copy=True) survives del image", arr[0, 0, 0] == 99.0) + del arr + + image = make_image(fill=77.0) + arr = itk.array_from_image(image) + del image + gc.collect() + check("itk.array_from_image survives del image", arr[0, 0, 0] == 77.0) + del arr + + gc.collect() + + # ============================================================== + # Section 2: __array__ zero-copy (3.10-3.11 primary path) + # ============================================================== + print() + print("=" * 60) + print("Section 2: __array__ zero-copy (NDArrayITKBase holds ref)") + print("=" * 60) + + if HAS_ZEROCOPY: + image = make_image(fill=33.0) + arr = image.__array__() + check("__array__() value matches", arr[0, 0, 0] == 33.0) + image.SetPixel([1, 1, 1], 123.0) + check("__array__() zero-copy: mutation visible", arr[1, 1, 1] == 123.0) + del image + gc.collect() + check( + "__array__() survives del image", + arr[0, 0, 0] == 33.0, + f"expected 33.0, got {arr[0, 0, 0]}", + ) + del arr + else: + skip("__array__() zero-copy tests", _SKIP_ZEROCOPY) + + image = make_image(fill=55.0) + arr = itk.array_view_from_image(image) + check("array_view_from_image value matches", arr[0, 0, 0] == 55.0) + del image + gc.collect() + check( + "array_view_from_image survives del image", + arr[0, 0, 0] == 55.0, + f"expected 55.0, got {arr[0, 0, 0]}", + ) + del arr + + gc.collect() + + # ============================================================== + # Section 3: np.asarray (critical path — all versions) + # ============================================================== + print() + print("=" * 60) + print("Section 3: np.asarray lifetime (critical path)") + print("=" * 60) + + if HAS_ZEROCOPY: + image = make_image(fill=11.0) + arr = np.asarray(image) + check("np.asarray value matches", arr[0, 0, 0] == 11.0) + image.SetPixel([0, 0, 0], 22.0) + check("np.asarray zero-copy: mutation visible", arr[0, 0, 0] == 22.0) + # THE CRITICAL TEST (thewtex's concern) + del image + gc.collect() + check( + "np.asarray survives del image", + arr[0, 0, 0] == 22.0, + f"expected 22.0, got {arr[0, 0, 0]}", + ) + del arr + + for ptype, fill_val, dtype_name in [ + (itk.UC, 200, "uint8"), + (itk.SS, -1000, "int16"), + (itk.F, 3.14, "float32"), + ]: + image = make_image( + pixel_type=ptype, dimension=2, size=[32, 32], fill=fill_val + ) + arr = np.asarray(image) + del image + gc.collect() + check( + f"np.asarray({dtype_name}) survives del image", + abs(float(arr[0, 0]) - fill_val) < 0.01, + f"expected {fill_val}, got {arr[0, 0]}", + ) + del arr + else: + skip("np.asarray zero-copy lifetime tests", _SKIP_ZEROCOPY) + + gc.collect() + + # ============================================================== + # Section 4: __buffer__ / memoryview (PEP 688) + # ============================================================== + print() + print("=" * 60) + print("Section 4: __buffer__ / memoryview lifetime") + print("=" * 60) + + if HAS_BUFFER: + image = make_image(pixel_type=itk.UC, dimension=2, size=[10, 10], fill=7) + mv = image.__buffer__() + check("__buffer__() value matches", mv[0, 0] == 7) + del image + gc.collect() + check( + "__buffer__() survives del image", + mv[0, 0] == 7, + f"expected 7, got {mv[0, 0]}", + ) + del mv + + if sys.version_info >= (3, 12): + image = make_image(pixel_type=itk.SS, dimension=2, size=[8, 8], fill=-5) + mv = memoryview(image) + check("memoryview(image) value matches", mv[0, 0] == -5) + del image + gc.collect() + check( + "memoryview(image) survives del image", + mv[0, 0] == -5, + f"expected -5, got {mv[0, 0]}", + ) + del mv + else: + skip( + "memoryview(image) auto", + f"requires Python 3.12+, have {sys.version_info[:2]}", + ) + + image = make_image(pixel_type=itk.F, dimension=3, size=[4, 5, 6], fill=2.5) + mv = image.__buffer__() + check("__buffer__() 3D float value", abs(mv[0, 0, 0] - 2.5) < 1e-5) + del image + gc.collect() + check( + "__buffer__() 3D float survives del image", + abs(mv[0, 0, 0] - 2.5) < 1e-5, + ) + del mv + else: + skip("__buffer__ lifetime tests", _SKIP_BUFFER) + + gc.collect() + + # ============================================================== + # Section 5: Chained references and GC stress + # ============================================================== + print() + print("=" * 60) + print("Section 5: Chained references and GC stress") + print("=" * 60) + + if HAS_ZEROCOPY: + # slice chain + image = make_image(fill=50.0) + arr = np.asarray(image) + sliced = arr[2:6, 2:6, 0:4] + check("slice value before del", sliced[0, 0, 0] == 50.0) + del image + gc.collect() + check("slice survives del image", sliced[0, 0, 0] == 50.0) + del arr + gc.collect() + check("slice survives del arr AND del image", sliced[0, 0, 0] == 50.0) + del sliced + + # transpose chain + image = make_image(fill=60.0) + arr = np.asarray(image) + transposed = arr.T + del image + gc.collect() + check("transposed view survives del image", transposed[0, 0, 0] == 60.0) + del arr, transposed + + # ravel chain + image = make_image(pixel_type=itk.UC, dimension=2, size=[8, 8], fill=9) + arr = np.asarray(image) + flat = arr.ravel() + del image + gc.collect() + check("raveled view survives del image", flat[0] == 9) + del arr, flat + + # multi-ref + image = make_image(fill=70.0) + arr1 = np.asarray(image) + arr2 = np.asarray(image) + if HAS_BUFFER: + mv = image.__buffer__() + del image + gc.collect() + check("arr1 survives (multi-ref)", arr1[0, 0, 0] == 70.0) + check("arr2 survives (multi-ref)", arr2[0, 0, 0] == 70.0) + if HAS_BUFFER: + check("memoryview survives (multi-ref)", abs(mv[0, 0, 0] - 70.0) < 1e-5) + del mv + del arr1, arr2 + + # rapid create-export-delete — crash-freedom is the assertion; + # any use-after-free would raise or segfault inside the loop. + for i in range(100): + img = make_image(pixel_type=itk.UC, dimension=2, size=[4, 4], fill=i % 256) + a = np.asarray(img) + del a, img + gc.collect() + check("100x create-export-delete cycle (no crash)", True) + else: + skip("chained reference tests", _SKIP_ZEROCOPY) + + # ============================================================== + # Section 6: Filter pipeline lifetime + # ============================================================== + print() + print("=" * 60) + print("Section 6: Filter pipeline output lifetime") + print("=" * 60) + + if HAS_ZEROCOPY: + input_image = make_image(fill=10.0) + median = itk.MedianImageFilter.New(input_image, Radius=1) + median.Update() + output = median.GetOutput() + arr = np.asarray(output) + check("filter output np.asarray value", abs(arr[2, 2, 2] - 10.0) < 1e-5) + del median, input_image, output + gc.collect() + check( + "filter output array survives pipeline deletion", + abs(arr[2, 2, 2] - 10.0) < 1e-5, + ) + del arr + + input_image = make_image(fill=20.0) + median = itk.MedianImageFilter.New(input_image, Radius=1) + median.Update() + arr = np.asarray(median.GetOutput()) + del median, input_image + gc.collect() + check( + "inline GetOutput().asarray survives", + abs(arr[2, 2, 2] - 20.0) < 1e-5, + ) + del arr + else: + skip("pipeline lifetime tests", _SKIP_ZEROCOPY) + + gc.collect() + + # ============================================================== + # Section 7: Weak reference GC verification + # ============================================================== + print() + print("=" * 60) + print("Section 7: Weak reference GC verification") + print("=" * 60) + + if HAS_ZEROCOPY: + # np.asarray path + image = make_image() + ref = weakref.ref(image) + arr = np.asarray(image) + del image + gc.collect() + check("image alive while array exists", ref() is not None) + del arr + gc.collect() + check( + "image GC'd after array deleted", + ref() is None, + "image leaked -- ref() still alive", + ) + else: + skip("np.asarray weak ref test", _SKIP_ZEROCOPY) + + if HAS_BUFFER: + # __buffer__ path + image = make_image() + ref = weakref.ref(image) + mv = image.__buffer__() + del image + gc.collect() + check("image alive while memoryview exists", ref() is not None) + del mv + gc.collect() + check( + "image GC'd after memoryview deleted", + ref() is None, + "image leaked via __buffer__ path", + ) + else: + skip("__buffer__ weak ref test", _SKIP_BUFFER) + + # array_view_from_image path + image = make_image() + ref = weakref.ref(image) + view = itk.array_view_from_image(image) + del image + gc.collect() + check("image alive while view exists", ref() is not None) + del view + gc.collect() + check( + "image GC'd after view deleted", + ref() is None, + "image leaked via array_view_from_image", + ) + + # deep copy must NOT hold image reference + image = make_image() + ref = weakref.ref(image) + arr_copy = itk.array_from_image(image) + del image + gc.collect() + check( + "image GC'd immediately after array_from_image", + ref() is None, + "deep copy is holding a reference to the image", + ) + del arr_copy + + image = make_image() + ref = weakref.ref(image) + arr_copy = np.array(image, copy=True) + del image + gc.collect() + check("image GC'd immediately after np.array(copy=True)", ref() is None) + del arr_copy + + if HAS_ARRAY_COPY_PARAM: + image = make_image() + ref = weakref.ref(image) + arr_copy = image.__array__(copy=True) + del image + gc.collect() + check("image GC'd immediately after __array__(copy=True)", ref() is None) + del arr_copy + else: + skip("__array__(copy=True) GC test", _SKIP_COPY) + + if HAS_ZEROCOPY: + # chained slice + image = make_image() + ref = weakref.ref(image) + arr = np.asarray(image) + sliced = arr[10:14, 10:14, 0:4] + del image, arr + gc.collect() + check("image alive while slice exists", ref() is not None) + del sliced + gc.collect() + check( + "image GC'd after slice deleted", + ref() is None, + "image leaked via slice chain", + ) + + # pipeline output + input_image = make_image(fill=10.0) + median = itk.MedianImageFilter.New(input_image, Radius=1) + median.Update() + output = median.GetOutput() + ref = weakref.ref(output) + arr = np.asarray(output) + del output, median, input_image + gc.collect() + check("pipeline output alive while array exists", ref() is not None) + del arr + gc.collect() + check( + "pipeline output GC'd after array deleted", + ref() is None, + "pipeline output leaked", + ) + else: + skip("chained/pipeline weak ref tests", _SKIP_ZEROCOPY) + + gc.collect() + + # ============================================================== + # Section 8: Reference count verification + # ============================================================== + print() + print("=" * 60) + print("Section 8: Reference count verification") + print("=" * 60) + + if HAS_ZEROCOPY: + image = make_image() + base_refcount = sys.getrefcount(image) + + arr = np.asarray(image) + after_asarray = sys.getrefcount(image) + check( + "np.asarray adds reference", + after_asarray > base_refcount, + f"base={base_refcount}, after={after_asarray}", + ) + + arr2 = np.asarray(image) + after_second = sys.getrefcount(image) + check( + "second np.asarray adds another reference", + after_second > after_asarray, + f"after_first={after_asarray}, after_second={after_second}", + ) + + del arr, arr2 + gc.collect() + restored = sys.getrefcount(image) + check( + "refcount restored after del arrays", + restored == base_refcount, + f"base={base_refcount}, restored={restored}", + ) + else: + skip("np.asarray refcount tests", _SKIP_ZEROCOPY) + image = make_image() + base_refcount = sys.getrefcount(image) + + if HAS_BUFFER: + mv = image.__buffer__() + after_buffer = sys.getrefcount(image) + check( + "__buffer__ adds reference", + after_buffer > base_refcount, + f"base={base_refcount}, after={after_buffer}", + ) + del mv + gc.collect() + restored2 = sys.getrefcount(image) + check( + "refcount restored after del memoryview", + restored2 == base_refcount, + f"base={base_refcount}, restored={restored2}", + ) + else: + skip("__buffer__ refcount tests", _SKIP_BUFFER) + + del image + gc.collect() + + # ============================================================== + # Section 9: Circular reference detection + # ============================================================== + print() + print("=" * 60) + print("Section 9: Circular reference detection") + print("=" * 60) + + image = make_image(size=[8, 8, 8]) + arr = np.asarray(image) + gc.collect() + garbage_before = len(gc.garbage) + del arr, image + gc.collect() + check( + "no uncollectable garbage from np.asarray", + len(gc.garbage) == garbage_before, + f"garbage grew from {garbage_before} to {len(gc.garbage)}", + ) + + if HAS_BUFFER: + image = make_image(size=[8, 8, 8]) + mv = image.__buffer__() + gc.collect() + garbage_before = len(gc.garbage) + del mv, image + gc.collect() + check( + "no uncollectable garbage from __buffer__", + len(gc.garbage) == garbage_before, + f"garbage grew from {garbage_before} to {len(gc.garbage)}", + ) + else: + skip("__buffer__ garbage test", _SKIP_BUFFER) + + # -- end of function: all locals go out of scope -- + + +# ================================================================== +# Main: run tests, then verify no RSS growth (leak detection) +# ================================================================== + +# Warm up ITK lazy loading before measuring +_warmup = make_image(size=[4, 4, 4]) +_ = np.asarray(_warmup) +del _warmup, _ +gc.collect() + +# Run all lifetime and GC tests inside a function scope +run_lifetime_and_leak_tests() +gc.collect() + +# After the function returns, all locals are out of scope. +# Verify no RSS growth from repeated create-export-delete cycles. +print() +print("=" * 60) +print("Section 10: RSS memory growth (leak detection)") +print("=" * 60) + +rss_available = get_rss_bytes() is not None +if rss_available: + N_ITERATIONS = 200 + IMAGE_SIZE = [64, 64, 64] + + export_paths = [ + ("np.asarray", lambda img: np.asarray(img)), + ("array_view_from_image", lambda img: itk.array_view_from_image(img)), + ("array_from_image", lambda img: itk.array_from_image(img)), + ] + if HAS_BUFFER: + export_paths.insert(1, ("__buffer__", lambda img: img.__buffer__())) + + for path_name, export_fn in export_paths: + gc.collect() + rss_before = get_rss_bytes() + for _ in range(N_ITERATIONS): + img = make_image(size=IMAGE_SIZE) + exported = export_fn(img) + del exported, img + gc.collect() + rss_after = get_rss_bytes() + growth_mb = (rss_after - rss_before) / (1024 * 1024) + check( + f"{path_name} RSS growth < 20 MB ({growth_mb:.1f} MB)", + growth_mb < 20, + f"{growth_mb:.1f} MB over {N_ITERATIONS} iterations", + ) +else: + skip("RSS growth tests", "RSS measurement not available on this platform") + + +# ================================================================== +# Summary +# ================================================================== +print() +print("=" * 60) +print(f"Python {sys.version_info[0]}.{sys.version_info[1]}.{sys.version_info[2]}") +print(f"NumPy {np.__version__}") +print( + f"Features: __buffer__={HAS_BUFFER}, copy_param={HAS_ARRAY_COPY_PARAM}, " + f"zero_copy={HAS_ZEROCOPY}" +) +print(f"TOTAL: {passed} passed, {failed} failed, {skipped} skipped") +print("=" * 60) + +if failed > 0: + print("LIFETIME TESTS FAILED!") + sys.exit(1) +else: + print("All lifetime and memory leak tests passed.") diff --git a/Modules/Core/Common/wrapping/test/itkImageTest.py b/Modules/Core/Common/wrapping/test/itkImageTest.py index 140806f119a..edd3c736c11 100644 --- a/Modules/Core/Common/wrapping/test/itkImageTest.py +++ b/Modules/Core/Common/wrapping/test/itkImageTest.py @@ -15,21 +15,50 @@ # limitations under the License. # # ========================================================================== +"""Test itk.Image numpy/buffer protocol integration. + +Tests are automatically skipped when features are not available in the +current ITK build (e.g., __buffer__ requires PEP 688 support added in +ITK 6.x). This allows cherry-picking into older ITK releases. +""" +import sys import itk import numpy as np -Dimension = 2 -PixelType = itk.UC +skipped = 0 + -ImageType = itk.Image[PixelType, Dimension] +def skip(name, reason): + global skipped + skipped += 1 + print(f" SKIP: {name} ({reason})") -image_size = [10, 10] -image = ImageType.New() -image.SetRegions(image_size) +# --- Feature detection --- +_probe = itk.Image[itk.UC, 2].New() +_probe.SetRegions([2, 2]) +_probe.Allocate() +HAS_BUFFER = hasattr(_probe, "__buffer__") +HAS_ARRAY_COPY_PARAM = False +try: + _probe.__array__(copy=None) + HAS_ARRAY_COPY_PARAM = True +except TypeError: + pass +# Detect zero-copy np.asarray +_probe.FillBuffer(1) +_a = np.asarray(_probe) +_probe.SetPixel([0, 0], 99) +HAS_ZEROCOPY = _a[0, 0] == 99 +del _a, _probe + +# --- Setup --- +image = itk.Image[itk.UC, 2].New() +image.SetRegions([10, 10]) image.Allocate() image.FillBuffer(4) +# --- Basic __array__ (all ITK versions) --- array = image.__array__() assert array[0, 0] == 4 assert array[0, 1] == 4 @@ -39,3 +68,99 @@ assert array[0, 0] == 4 assert array[0, 1] == 4 assert isinstance(array, np.ndarray) + +# --- PEP 688 buffer protocol tests --- +if HAS_BUFFER: + 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 + else: + skip( + "memoryview(image) auto", + f"requires Python 3.12+, have {sys.version_info[:2]}", + ) +else: + skip("PEP 688 __buffer__ tests", "__buffer__ not available in this ITK version") + +# --- np.asarray zero-copy --- +if HAS_ZEROCOPY: + arr_view = np.asarray(image) + assert arr_view.shape == (10, 10) + assert arr_view.dtype == np.uint8 + image.SetPixel([1, 2], 99) + assert arr_view[2, 1] == 99 # zero-copy verified + + arr_a = image.__array__() + image.SetPixel([0, 0], 77) + assert arr_a[0, 0] == 77 # zero-copy verified +else: + skip( + "np.asarray zero-copy tests", + "zero-copy np.asarray not available in this ITK version", + ) + +# --- __array__ copy parameter (NumPy 2.0 protocol) --- +if HAS_ARRAY_COPY_PARAM: + arr_none = image.__array__(copy=None) + image.SetPixel([0, 0], 55) + assert arr_none[0, 0] == 55 # zero-copy + + arr_true = image.__array__(copy=True) + image.SetPixel([0, 0], 66) + assert arr_true[0, 0] == 55 # copy: still old value + + arr_false = image.__array__(copy=False) + image.SetPixel([0, 0], 88) + assert arr_false[0, 0] == 88 # zero-copy + + try: + _ = image.__array__(dtype=np.float64, copy=False) + assert False, "copy=False with dtype conversion should raise ValueError" + except ValueError: + pass # expected: dtype conversion requires a copy + + arr_same_dtype = image.__array__(dtype=np.uint8, copy=False) + image.SetPixel([0, 0], 91) + assert arr_same_dtype[0, 0] == 91 # zero-copy, no dtype conversion +else: + skip( + "__array__(copy=...) tests", + "__array__ copy parameter not available in this ITK version", + ) + +if skipped: + print(f"({skipped} test group(s) skipped)") +print("All buffer protocol tests passed.")