Using ITK in a task type

ITK can be used in a task type via SimpleITK if Core.Python.SimpleITK is installed.

For example Example.Itk.AnisotropicDiffusion uses SimpleITK.CurvatureAnisotropicDiffusionImageFilter:

import omni
import xfworkerutil

import SimpleITK


worker = xfworkerutil.XFWorkerJob()

inputs = worker.job['inputs']
iterations = inputs['iterations']
time_step = inputs['time_step']
conductance = inputs['conductance']
method = inputs['method']
convert_back_from_float = inputs['convert_back_from_float']

# Download input files (if not local)
image = worker.get_omni_image(inputs['image'])

# Use ITK for image processing
itk_image = SimpleITK.GetImageFromArray(image.pixels)

anisotropicFilter = SimpleITK.CurvatureAnisotropicDiffusionImageFilter() if method == 'curvature' else SimpleITK.GradientAnisotropicDiffusionImageFilter()
anisotropicFilter.SetNumberOfIterations(iterations)
anisotropicFilter.SetTimeStep(time_step)
anisotropicFilter.SetConductanceParameter(conductance)

# Cast the original_image to floating-point
itk_image = SimpleITK.Cast(itk_image, SimpleITK.sitkFloat32)
itk_image = anisotropicFilter.Execute(itk_image)
itk_image_pixels = SimpleITK.GetArrayFromImage(itk_image)

output_pixel_type = omni.Image.PIXEL_TYPE_BY_NUMPY_DTYPE[itk_image_pixels.dtype]
if convert_back_from_float:
    output_pixel_type = image.pixel_type
    itk_image_pixels = itk_image_pixels.astype(omni.Image.PIXEL_TYPE_TO_NUMPY_DTYPE[output_pixel_type])

# Save the image
result_omni_path = worker.create_related_path(image.metadata_path, '_itk_diff.omni')
result_raw_path = worker.create_related_path(image.metadata_path, '_itk_diff.raw')
itk_image_pixels.tofile(str(result_raw_path))

result_image = image.save_derived_image(result_omni_path, result_raw_path, {
    'pixel_type': output_pixel_type,
})

# Upload output files
result_image_output = worker.upload_omni_output(result_image)

# Save outputs
outputs = {
    "image": result_image_output,
}
worker.finish(outputs)