Skip to content

Commit

Permalink
BUG: Support serialization of an empty itk.Image
Browse files Browse the repository at this point in the history
If an itk.Image's memory has not been allocated return None from
itk.array_view_from_image.

Pickle and unpickles an unallocated itk.Image to support sending the
image's metadata with Dask. #4267.
  • Loading branch information
thewtex committed Oct 27, 2023
1 parent a661e0e commit 0a08d16
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 2 deletions.
6 changes: 6 additions & 0 deletions Modules/Bridge/NumPy/wrapping/PyBuffer.i.in
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@
i.e. k,j,i versus i,j,k. However C-order indexing is expected by most
algorithms in NumPy / SciPy.
"""

if image.GetBufferPointer() is None:
return None

if update:
# Ensure the image regions and image pixel buffer have been updated
# correctly
Expand Down Expand Up @@ -46,6 +50,8 @@

This is a deep copy of the image buffer and is completely safe and without potential side effects.
"""
if image.GetBufferPointer() is None:
return None

arrayView = itkPyBuffer@[email protected](image, keep_axes, update)

Expand Down
32 changes: 32 additions & 0 deletions Modules/Bridge/NumPy/wrapping/test/itkPyBufferTest.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,38 @@ def test_NDArrayITKBase_pickle(self):
equal = (recon_obj == ndarray_itk_base).all()
assert equal, "Different results before and after pickle"

def test_EmptyImage_pickle(self):
"""
Test the serialization of an empty itk.Image
"""
Dimension = 3
ImageType = itk.Image[itk.UC, Dimension]
RegionType = itk.ImageRegion[Dimension]

region = RegionType()
region.SetSize(0, 6)
region.SetSize(1, 6)
region.SetSize(2, 6)

image = ImageType.New()
image.SetRegions(region)
spacing = [3.0, 4.0, 5.0]
image.SetSpacing(spacing)
# Before allocation
# scalarImage.Allocate(True)

import pickle
pickled = pickle.dumps(image)
reloaded = pickle.loads(pickled)
reloaded_spacing = reloaded.GetSpacing()
spacing = reloaded.GetSpacing()
assert spacing[0] == reloaded_spacing[0]
assert spacing[1] == reloaded_spacing[1]
assert spacing[2] == reloaded_spacing[2]

ndarray_itk_base = itk.array_view_from_image(image)
assert ndarray_itk_base == None

def test_NumPyBridge_itkScalarImage(self):
"Try to convert all pixel types to NumPy array view"

Expand Down
9 changes: 7 additions & 2 deletions Wrapping/Generators/Python/itk/support/extras.py
Original file line number Diff line number Diff line change
Expand Up @@ -805,12 +805,17 @@ def image_from_dict(image_dict: Dict) -> "itkt.Image":
import itk

ImageType = image_type_from_wasm_type(image_dict["imageType"])
image = itk.PyBuffer[ImageType].GetImageViewFromArray(image_dict["data"])
if image_dict["data"] is None:
image = ImageType.New()
image.SetRegions(image_dict["size"])
image.Allocate(True)
else:
image = itk.PyBuffer[ImageType].GetImageViewFromArray(image_dict["data"])
image.SetRegions(image_dict["size"])
image.SetOrigin(image_dict["origin"])
image.SetSpacing(image_dict["spacing"])
image.SetDirection(image_dict["direction"])
image.SetObjectName(image_dict["name"])
image.SetRegions(image_dict["size"])
return image


Expand Down

0 comments on commit 0a08d16

Please sign in to comment.