Error when multiprocessing coregistration

[plaresmedima]

Hi,

We are trying to parallelize coregistrations of large time series of images, but we are running into an error when we pass parameter objects to the functions.

Here is a simple code example that reproduces the issue:

import os
import multiprocessing
import numpy as np
import itk

# Define a function that coregisters two images with a given parameter object:
def coregister(args):
    fixed, moving, param_obj = args
    coregistered, _ = itk.elastix_registration_method(
        itk.GetImageFromArray(fixed), 
        itk.GetImageFromArray(moving), 
        parameter_object=param_obj)
    return itk.GetArrayFromImage(coregistered)


if __name__ == "__main__":

    # Create two time series
    nt = 2  # Note in real-world examples this would be in the 100's
    fixed = np.zeros((128,128,nt), dtype=np.float32)
    moving = np.zeros((128,128,nt), dtype=np.float32)

    # Generate a parameter object
    param_obj = itk.ParameterObject.New()
    bspline = param_obj.GetDefaultParameterMap('bspline')
    param_obj.AddParameterMap(bspline) 

    # Build a list of arguments
    args = [(fixed[:,:,t], moving[:,:,t], param_obj) for t in range(nt)]

    # Process them in parallel using the coregister function
    pool = multiprocessing.Pool(processes=os.cpu_count())
    coreg = pool.map(coregister, args)

This generates a type error:

TypeError: cannot pickle 'SwigPyObject' object

A workaround is to pass the elastix parameters to the coregister() function as a dictionary, and then generate the parameter object inside the coregister() function. However, the function itk.ParameterObject.New() is so slow that this completely offsets any benefits of the parallellization.

Do you have a better solution? Many thanks for any help you are able to give..

[plaresmedima]

Going through the list I suppose this is the same issue as highlighted in

#257