Skip to content

Commit

Permalink
BUG: Support passing a single transform to transform write
Browse files Browse the repository at this point in the history
This is a common use case. Without this patch, a segfault results.
  • Loading branch information
thewtex committed May 3, 2024
1 parent afc4879 commit f22076b
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 3 deletions.
3 changes: 3 additions & 0 deletions Wrapping/Generators/Python/Tests/extras.py
Expand Up @@ -348,6 +348,9 @@ def custom_callback(name, progress):
transform_dict = itk.dict_from_transform(transforms)
transform_back = itk.transform_from_dict(transform_dict)

# Write single transform
itk.transformwrite(transforms[0], sys.argv[7], compression=True)

# pipeline, auto_pipeline and templated class are tested in other files

# BridgeNumPy
Expand Down
8 changes: 5 additions & 3 deletions Wrapping/Generators/Python/itk/support/extras.py
Expand Up @@ -1453,7 +1453,7 @@ def transformread(filename: fileiotype) -> List["itkt.TransformBase"]:


def transformwrite(
transforms: List["itkt.TransformBase"],
transforms: Union["itkt.TransformBase", List["itkt.TransformBase"]],
filename: fileiotype,
compression: bool = False,
) -> None:
Expand All @@ -1462,8 +1462,8 @@ def transformwrite(
Parameters
----------
transforms: list of itk.TransformBaseTemplate[itk.D]
Python list of the transforms to write.
transforms: single transform or list of transforms
Single transform Python list of the itk.TransformBaseTemplate[itk.D]'s to write.
filename:
Path to the transform file (typically a .h5 file).
Expand All @@ -1476,6 +1476,8 @@ def transformwrite(
writer = itk.TransformFileWriterTemplate[itk.D].New()
writer.SetFileName(f"{filename}")
writer.SetUseCompression(compression)
if not isinstance(transforms, list):
transforms = [transforms]
for transform in transforms:
writer.AddTransform(transform)
writer.Update()
Expand Down

0 comments on commit f22076b

Please sign in to comment.