Skip to content

Commit

Permalink
ENH: wasm pipeline input parameters, initial transform support
Browse files Browse the repository at this point in the history
  • Loading branch information
@thewtex
thewtex committed Sep 25, 2023
1 parent 5ee4535 commit 6abb5c3
Show file tree
Hide file tree
Showing 15 changed files with 275 additions and 15 deletions.
Binary file modified examples/exampleoutput/itk_composite_transform.h5
View file
Binary file not shown.
31 changes: 31 additions & 0 deletions wasm/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,46 @@ target_link_libraries(write-parameter-files PUBLIC ${ITK_LIBRARIES})
add_executable(default-parameter-map default-parameter-map.cxx)
target_link_libraries(default-parameter-map PUBLIC ${ITK_LIBRARIES})


enable_testing()

add_test(NAME elastix-wasm-test
COMMAND elastix
${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/parameters_single.json
${CMAKE_CURRENT_BINARY_DIR}/CT_3D_lung_registered.iwi.cbor
${CMAKE_CURRENT_BINARY_DIR}/CT_3D_lung.h5
-f ${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/CT_3D_lung_fixed.iwi.cbor
-m ${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/CT_3D_lung_moving.iwi.cbor
)

add_test(NAME elastix-wasm-2d-test
COMMAND elastix
${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/parameters_single.json
${CMAKE_CURRENT_BINARY_DIR}/CT_2D_head_registered.iwi.cbor
${CMAKE_CURRENT_BINARY_DIR}/CT_2D_head.h5
-f ${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/CT_2D_head_fixed.iwi.cbor
-m ${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/CT_2D_head_moving.iwi.cbor
)

add_test(NAME elastix-wasm-2d-initial-test
COMMAND elastix
${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/parameters_single.json
${CMAKE_CURRENT_BINARY_DIR}/CT_2D_head_registered.iwi.cbor
${CMAKE_CURRENT_BINARY_DIR}/CT_2D_head_initial.h5
-f ${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/CT_2D_head_fixed.iwi.cbor
-m ${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/CT_2D_head_moving.iwi.cbor
-i ${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/CT_2D_head_translation.h5
)

add_test(NAME elastix-wasm-2d-multiple-test
COMMAND elastix
${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/parameters_multiple.json
${CMAKE_CURRENT_BINARY_DIR}/CT_2D_head_multiple_registered.iwi.cbor
${CMAKE_CURRENT_BINARY_DIR}/CT_2D_head_multiple.h5
-f ${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/CT_2D_head_fixed.iwi.cbor
-m ${CMAKE_CURRENT_SOURCE_DIR}/test/data/input/CT_2D_head_moving.iwi.cbor
)

add_test(NAME read-parameter-files-test
COMMAND read-parameter-files
${CMAKE_CURRENT_BINARY_DIR}/parameters_single.json
Expand Down
133 changes: 123 additions & 10 deletions wasm/elastix-wasm.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,22 @@
#include "itkPipeline.h"
#include "itkInputImage.h"
#include "itkOutputImage.h"
#include "itkInputTextStream.h"
#include "itkSupportInputImageTypes.h"

#include "itkImage.h"
#include "itkTransformFileWriter.h"
#include "itkTransformFileReader.h"
#include "itkIdentityTransform.h"
#include "itkCompositeTransform.h"
#include "itkCompositeTransformIOHelper.h"
#include "itkCastImageFilter.h"

#include "rapidjson/document.h"
#include "rapidjson/prettywriter.h"
#include "rapidjson/stringbuffer.h"

#include <sstream>

template <typename TImage>
class PipelineFunctor
Expand All @@ -33,6 +44,7 @@ class PipelineFunctor
operator()(itk::wasm::Pipeline & pipeline)
{
using ImageType = TImage;
using ParametersValueType = double;

using InputImageType = itk::wasm::InputImage<ImageType>;
InputImageType fixedImage;
Expand All @@ -41,6 +53,14 @@ class PipelineFunctor
InputImageType movingImage;
pipeline.add_option("-m,--moving", movingImage, "Moving image")->type_name("INPUT_IMAGE");

std::string initialTransformFile;
pipeline.add_option("-i,--initial-transform", initialTransformFile, "Initial transform to apply before registrtion ")->type_name("INPUT_BINARY_FILE");

itk::wasm::InputTextStream parameterObjectJson;
pipeline.add_option("parameter-object", parameterObjectJson, "Elastix parameter object representation")
->required()
->type_name("INPUT_JSON");

using OutputImageType = itk::wasm::OutputImage<ImageType>;
OutputImageType resultImage;
pipeline.add_option("result", resultImage, "Resampled moving image")->required()->type_name("OUTPUT_IMAGE");
Expand All @@ -52,18 +72,100 @@ class PipelineFunctor

ITK_WASM_PARSE(pipeline);

using RegistrationType = itk::ElastixRegistrationMethod<ImageType, ImageType>;
using FloatImageType = itk::Image<float, ImageType::ImageDimension>;
using CasterType = itk::CastImageFilter<ImageType, FloatImageType>;

typename CasterType::Pointer fixedCaster = CasterType::New();
fixedCaster->SetInput(fixedImage.Get());
ITK_WASM_CATCH_EXCEPTION(pipeline, fixedCaster->Update());

typename CasterType::Pointer movingCaster = CasterType::New();
movingCaster->SetInput(movingImage.Get());
ITK_WASM_CATCH_EXCEPTION(pipeline, movingCaster->Update());

using RegistrationType = itk::ElastixRegistrationMethod<FloatImageType, FloatImageType>;
typename RegistrationType::Pointer registration = RegistrationType::New();

rapidjson::Document document;
std::stringstream ss;
ss << parameterObjectJson.Get().rdbuf();
document.Parse(ss.str().c_str());

using ParameterObjectType = elastix::ParameterObject;
const auto parameterObject = ParameterObjectType::New();
const auto numParameterMaps = document.Size();
using ParameterMapType = std::map<std::string, std::vector<std::string>>;
std::vector<ParameterMapType> parameterMaps;
for (unsigned int i = 0; i < numParameterMaps; ++i)
{
const auto & parameterMapJson = document[i];
ParameterMapType parameterMap;
for (auto it = parameterMapJson.MemberBegin(); it != parameterMapJson.MemberEnd(); ++it)
{
const auto & key = it->name.GetString();
const auto & valueJson = it->value;
std::vector<std::string> value;
for (auto it2 = valueJson.Begin(); it2 != valueJson.End(); ++it2)
{
const auto & valueElement = it2->GetString();
value.push_back(valueElement);
}
parameterMap[key] = value;
}
parameterObject->AddParameterMap(parameterMap);
}

auto fixed = const_cast<ImageType *>(fixedImage.Get());
auto moving = const_cast<ImageType *>(movingImage.Get());
registration->SetFixedImage(fixed);
registration->SetMovingImage(moving);
registration->SetFixedImage(fixedCaster->GetOutput());
registration->SetMovingImage(movingCaster->GetOutput());
registration->SetParameterObject(parameterObject);

typename RegistrationType::TransformType::Pointer initialTransform;
using CompositeTransformType = itk::CompositeTransform<ParametersValueType, ImageType::ImageDimension>;
using CompositeHelperType = itk::CompositeTransformIOHelperTemplate<ParametersValueType>;
using TransformReaderType = itk::TransformFileReaderTemplate<ParametersValueType>;
typename TransformReaderType::Pointer transformReader = TransformReaderType::New();
if (!initialTransformFile.empty())
{
transformReader->SetFileName(initialTransformFile);
ITK_WASM_CATCH_EXCEPTION(pipeline, transformReader->Update());

if (transformReader->GetTransformList()->size() == 1)
{
auto firstTransform = transformReader->GetModifiableTransformList()->front();
if (!strcmp(firstTransform->GetNameOfClass(), "CompositeTransform"))
{
initialTransform = static_cast<CompositeTransformType *>(firstTransform.GetPointer());
registration->SetExternalInitialTransform(initialTransform);
}
// We could add support for other initial transform types here
else
{
std::cerr << "Initial transform is not a composite transform, which is not currently supported." << std::endl;
return EXIT_FAILURE;
}
}
else if (transformReader->GetTransformList()->size() > 1)
{
CompositeHelperType helper;
typename CompositeTransformType::Pointer compositeTransform = CompositeTransformType::New();
helper.SetTransformList(compositeTransform, *transformReader->GetModifiableTransformList());
initialTransform = compositeTransform;
registration->SetExternalInitialTransform(initialTransform);
}
}


ITK_WASM_CATCH_EXCEPTION(pipeline, registration->Update());

typename ImageType::Pointer outputImage = registration->GetOutput();
resultImage.Set(outputImage);
typename FloatImageType::Pointer outputImage = registration->GetOutput();
using ResultCasterType = itk::CastImageFilter<FloatImageType, ImageType>;
typename ResultCasterType::Pointer resultCaster = ResultCasterType::New();
resultCaster->SetInput(outputImage);
ITK_WASM_CATCH_EXCEPTION(pipeline, resultCaster->Update());
typename ImageType::ConstPointer result = resultCaster->GetOutput();
resultImage.Set(result);

const auto writer = itk::TransformFileWriter::New();

Expand All @@ -75,14 +177,25 @@ class PipelineFunctor
writer->SetFileName(outputTransform);
ITK_WASM_CATCH_EXCEPTION(pipeline, writer->Update());
}
// Reasonable to enable once we support injecting as an initial transform
// else if (!initialTransform.GetPointer() && registration->GetNumberOfTransforms() == 1)
// {
// auto transform = registration->GetNthTransform(0);
// typename RegistrationType::TransformType::ConstPointer registeredTransform =
// registration->ConvertToItkTransform(*transform);
// writer->SetInput(registeredTransform);
// writer->SetFileName(outputTransform);
// ITK_WASM_CATCH_EXCEPTION(pipeline, writer->Update());
// }
else
{
typename RegistrationType::TransformType::ConstPointer combinationTransform =
registration->GetCombinationTransform();
combinationTransform->Print(std::cout);
typename RegistrationType::TransformType::ConstPointer compositeTransform =
registration->ConvertToItkTransform(*combinationTransform);
writer->SetInput(compositeTransform);
typename CompositeTransformType::Pointer registeredCompositeTransform =
static_cast<CompositeTransformType *>(registration->ConvertToItkTransform(*combinationTransform).GetPointer());
registeredCompositeTransform->FlattenTransformQueue();
registeredCompositeTransform->SetAllTransformsToOptimizeOff();
writer->SetInput(registeredCompositeTransform);
writer->SetFileName(outputTransform);
ITK_WASM_CATCH_EXCEPTION(pipeline, writer->Update());
}
Expand All @@ -97,5 +210,5 @@ main(int argc, char * argv[])
itk::wasm::Pipeline pipeline("elastix", "Rigid and non-rigid registration of images.", argc, argv);

return itk::wasm::SupportInputImageTypes<PipelineFunctor, uint8_t, uint16_t, int16_t, double, float>::
Dimensions<2U, 3U, 4U>("-f,--fixed", pipeline);
Dimensions<2U, 3U>("-f,--fixed", pipeline);
}
12 changes: 11 additions & 1 deletion wasm/python/itkwasm-elastix-emscripten/itkwasm_elastix_emscripten/elastix_async.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,17 +20,25 @@
)

async def elastix_async(
parameter_object: Any,
fixed: Optional[Image] = None,
moving: Optional[Image] = None,
initial_transform: Optional[os.PathLike] = None,
) -> Tuple[Image, os.PathLike]:
"""Rigid and non-rigid registration of images.
:param parameter_object: Elastix parameter object representation
:type parameter_object: Any
:param fixed: Fixed image
:type fixed: Image
:param moving: Moving image
:type moving: Image
:param initial_transform: Initial transform to apply before registrtion
:type initial_transform: os.PathLike
:return: Resampled moving image
:rtype: Image
Expand All @@ -45,8 +53,10 @@ async def elastix_async(
kwargs["fixed"] = to_js(fixed)
if moving is not None:
kwargs["moving"] = to_js(moving)
if initial_transform is not None:
kwargs["initialTransform"] = to_js(BinaryFile(initial_transform))

outputs = await js_module.elastix(web_worker, **kwargs)
outputs = await js_module.elastix(web_worker, to_js(parameter_object), **kwargs)

output_web_worker = None
output_list = []
Expand Down
16 changes: 16 additions & 0 deletions wasm/python/itkwasm-elastix-wasi/itkwasm_elastix_wasi/elastix.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,17 +20,25 @@
)

def elastix(
parameter_object: Any,
fixed: Optional[Image] = None,
moving: Optional[Image] = None,
initial_transform: Optional[os.PathLike] = None,
) -> Tuple[Image, os.PathLike]:
"""Rigid and non-rigid registration of images.
:param parameter_object: Elastix parameter object representation
:type parameter_object: Any
:param fixed: Fixed image
:type fixed: Image
:param moving: Moving image
:type moving: Image
:param initial_transform: Initial transform to apply before registrtion
:type initial_transform: os.PathLike
:return: Resampled moving image
:rtype: Image
Expand All @@ -47,10 +55,12 @@ def elastix(
]

pipeline_inputs: List[PipelineInput] = [
PipelineInput(InterfaceTypes.JsonCompatible, parameter_object),
]

args: List[str] = ['--memory-io',]
# Inputs
args.append('0')
# Outputs
args.append('0')
args.append(str(PurePosixPath(transform)))
Expand All @@ -67,6 +77,12 @@ def elastix(
args.append('--moving')
args.append(input_count_string)

if initial_transform is not None:
input_file = str(PurePosixPath(initial_transform))
pipeline_inputs.append(PipelineInput(InterfaceTypes.BinaryFile, BinaryFile(initial_transform)))
args.append('--initial-transform')
args.append(input_file)


outputs = _pipeline.run(args, pipeline_outputs, pipeline_inputs)

Expand Down
Binary file modified wasm/python/itkwasm-elastix-wasi/itkwasm_elastix_wasi/wasm_modules/elastix.wasi.wasm
View file
Binary file not shown.
10 changes: 9 additions & 1 deletion wasm/python/itkwasm-elastix/itkwasm_elastix/elastix.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,23 +10,31 @@
)

def elastix(
parameter_object: Any,
fixed: Optional[Image] = None,
moving: Optional[Image] = None,
initial_transform: Optional[os.PathLike] = None,
) -> Tuple[Image, os.PathLike]:
"""Rigid and non-rigid registration of images.
:param parameter_object: Elastix parameter object representation
:type parameter_object: Any
:param fixed: Fixed image
:type fixed: Image
:param moving: Moving image
:type moving: Image
:param initial_transform: Initial transform to apply before registrtion
:type initial_transform: os.PathLike
:return: Resampled moving image
:rtype: Image
:return: Fixed-to-moving transform
:rtype: os.PathLike
"""
func = environment_dispatch("itkwasm_elastix", "elastix")
output = func(fixed=fixed, moving=moving)
output = func(parameter_object, fixed=fixed, moving=moving, initial_transform=initial_transform)
return output
10 changes: 9 additions & 1 deletion wasm/python/itkwasm-elastix/itkwasm_elastix/elastix_async.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,23 +10,31 @@
)

async def elastix_async(
parameter_object: Any,
fixed: Optional[Image] = None,
moving: Optional[Image] = None,
initial_transform: Optional[os.PathLike] = None,
) -> Tuple[Image, os.PathLike]:
"""Rigid and non-rigid registration of images.
:param parameter_object: Elastix parameter object representation
:type parameter_object: Any
:param fixed: Fixed image
:type fixed: Image
:param moving: Moving image
:type moving: Image
:param initial_transform: Initial transform to apply before registrtion
:type initial_transform: os.PathLike
:return: Resampled moving image
:rtype: Image
:return: Fixed-to-moving transform
:rtype: os.PathLike
"""
func = environment_dispatch("itkwasm_elastix", "elastix_async")
output = await func(fixed=fixed, moving=moving)
output = await func(parameter_object, fixed=fixed, moving=moving, initial_transform=initial_transform)
return output
Binary file modified wasm/typescript/dist/pipelines/elastix.wasm
View file
Binary file not shown.
Loading

0 comments on commit 6abb5c3

Please sign in to comment.