InsightSoftwareConsortium / ITKElastix

An ITK Python interface to elastix, a toolbox for rigid and nonrigid registration of images
Apache License 2.0
205 stars 22 forks source link

Make `itk.ParameterObject` and parameter maps serializable #257

Open tbirdso opened 11 months ago

tbirdso commented 11 months ago

Background

itk.ParameterObject acts as an ordered collection of Elastix parameter map dictionaries. We should make these C++ wrapper objects serializable to facilitate interactions with dask.distributed and other cases where registration may occur in a worker other than where Elastix parameters are originally specified.

Steps to Reproduce

import pickle
import itk

parameter_object = itk.ParameterObject.New()
parameter_object.AddParameterMap(itk.ParameterObject.GetDefaultParameterMap('rigid'))
bytestring = pickle.dumps(parameter_object)

Expected behavior

Parameter object is serialized and can be reconstructed with pickle.loads(bytestring)

Observed behavior

>>> pickle.dumps(po)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: cannot pickle 'SwigPyObject' object

Additional Notes

As a minimum effort it may be sufficient to package itk-elastix with Python helper methods to convert parameter maps and objects to and from Python types such as dict and list that already support serialization.

For a complete implementation we should wrap itk.ParameterObject and itk.elxParameterObjectPython.mapstringvectorstring with __setstate__ and __getstate__ methods. See Python's pickle documentation.

Related to work in https://github.com/InsightSoftwareConsortium/itk-dreg (cc @thewtex )

dzenanz commented 11 months ago

Could/should WriteParameterFile and related methods be used for this?

tbirdso commented 11 months ago

Could/should WriteParameterFile and related methods be used for this?

Possibly. Given that Elastix parameter maps nearly implement the Python dict interface it may be easier to go that route for in-memory serialization, rather than updating WriteParameterFile to write to an in-memory location.

For instance, to convert a single parameter map to a pickleable Python dict:

>>> dict(po.GetDefaultParameterMap('rigid'))
{'AutomaticParameterEstimation': ('true',), 'AutomaticScalesEstimation': ('true',), 'CheckNumberOfSamples': ('true',), 'DefaultPixelValue': ('0',), 'FinalBSplineInterpolationOrder': ('3',), 'FixedImagePyramid': ('FixedSmoothingImagePyramid',), 'ImageSampler': ('RandomCoordinate',), 'Interpolator': ('LinearInterpolator',), 'MaximumNumberOfIterations': ('256',), 'MaximumNumberOfSamplingAttempts': ('8',), 'Metric': ('AdvancedMattesMutualInformation',), 'MovingImagePyramid': ('MovingSmoothingImagePyramid',), 'NewSamplesEveryIteration': ('true',), 'NumberOfResolutions': ('4',), 'NumberOfSamplesForExactGradient': ('4096',), 'NumberOfSpatialSamples': ('2048',), 'Optimizer': ('AdaptiveStochasticGradientDescent',), 'Registration': ('MultiResolutionRegistration',), 'ResampleInterpolator': ('FinalBSplineInterpolator',), 'Resampler': ('DefaultResampler',), 'ResultImageFormat': ('nii',), 'Transform': ('EulerTransform',), 'WriteIterationInfo': ('false',), 'WriteResultImage': ('true',)}

To convert a parameter object to a pickleable list of Python dicts:

>>> [dict(parameter_object.GetParameterMap(map_index)) for map_index in range(parameter_object.GetNumberOfParameterMaps())]
[{'AutomaticParameterEstimation': ('true',), 'AutomaticScalesEstimation': ('true',), 'CheckNumberOfSamples': ('true',), 'DefaultPixelValue': ('0',), 'FinalBSplineInterpolationOrder': ('3',), 'FixedImagePyramid': ('FixedSmoothingImagePyramid',), 'ImageSampler': ('RandomCoordinate',), 'Interpolator': ('LinearInterpolator',), 'MaximumNumberOfIterations': ('256',), 'MaximumNumberOfSamplingAttempts': ('8',), 'Metric': ('AdvancedMattesMutualInformation',), 'MovingImagePyramid': ('MovingSmoothingImagePyramid',), 'NewSamplesEveryIteration': ('true',), 'NumberOfResolutions': ('4',), 'NumberOfSamplesForExactGradient': ('4096',), 'NumberOfSpatialSamples': ('2048',), 'Optimizer': ('AdaptiveStochasticGradientDescent',), 'Registration': ('MultiResolutionRegistration',), 'ResampleInterpolator': ('FinalBSplineInterpolator',), 'Resampler': ('DefaultResampler',), 'ResultImageFormat': ('nii',), 'Transform': ('EulerTransform',), 'WriteIterationInfo': ('false',), 'WriteResultImage': ('true',)}]
thewtex commented 11 months ago

To convert a parameter object to a pickleable list of Python dicts:

We can create these function in itk-dreg for now, apply before returning and when entering the block registration method, and integrate into itk-elastix as a later step.

tbirdso commented 11 months ago

To convert a parameter object to a pickleable list of Python dicts:

We can create these function in itk-dreg for now, apply before returning and when entering the block registration method, and integrate into itk-elastix as a later step.

Agreed, this issue is not holding up itk-dreg development. It will be helpful to revisit in the long term.

thewtex commented 5 months ago

With #278, it is much easier to add pure python modules to the package.