From 3777b45a25ad0f504bcb12b8c5a3bdad1ec64c74 Mon Sep 17 00:00:00 2001
From: Matt McCormick <matt.mccormick@kitware.com>
Date: Thu, 2 May 2024 13:10:39 -0400
Subject: [PATCH] ENH: Add dict_from_polyline, polyline_from_dict

Support numpy-based and Python dictionary representation of
itk.PolyLineParametricPath.

Closes #3278
---
 Wrapping/Generators/Python/PyBase/pyBase.i    | 44 +++++++++++++++++++
 Wrapping/Generators/Python/Tests/extras.py    | 14 ++++++
 .../Generators/Python/itk/support/extras.py   | 26 +++++++++++
 Wrapping/TypedefMacros.cmake                  |  4 ++
 4 files changed, 88 insertions(+)

diff --git a/Wrapping/Generators/Python/PyBase/pyBase.i b/Wrapping/Generators/Python/PyBase/pyBase.i
index 4450284b7c0..76511a23402 100644
--- a/Wrapping/Generators/Python/PyBase/pyBase.i
+++ b/Wrapping/Generators/Python/PyBase/pyBase.i
@@ -735,6 +735,50 @@ str = str
 
 %enddef
 
+%define DECL_PYTHON_POLYLINEPARAMETRICPATH_CLASS(swig_name)
+    %extend swig_name {
+        %pythoncode %{
+            def keys(self):
+                """
+                Return keys related to the polyline's metadata.
+                These keys are used in the dictionary resulting from dict(polyline).
+                """
+                result = ['name', 'vertexList']
+                return result
+
+            def __getitem__(self, key):
+                """Access metadata keys, see help(pointset.keys), for string keys."""
+                import itk
+                if isinstance(key, str):
+                    state = itk.dict_from_polyline(self)
+                    return state[key]
+
+            def __setitem__(self, key, value):
+                if isinstance(key, str):
+                    import numpy as np
+                    if key == 'name':
+                        self.SetObjectName(value)
+                    elif key == 'vertexList':
+                        self.GetVertexList().Initialize()
+                        for vertex in value:
+                            polyline.AddVertex(vertex)
+
+            def __getstate__(self):
+                """Get object state, necessary for serialization with pickle."""
+                import itk
+                state = itk.dict_from_polyline(self)
+                return state
+
+            def __setstate__(self, state):
+                """Set object state, necessary for serialization with pickle."""
+                import itk
+                deserialized = itk.polyline_from_dict(state)
+                self.__dict__['this'] = deserialized
+            %}
+    }
+
+%enddef
+
 %define DECL_PYTHON_MESH_CLASS(swig_name)
     %extend swig_name {
         %pythoncode %{
diff --git a/Wrapping/Generators/Python/Tests/extras.py b/Wrapping/Generators/Python/Tests/extras.py
index 773ffac149b..d3940fb6832 100644
--- a/Wrapping/Generators/Python/Tests/extras.py
+++ b/Wrapping/Generators/Python/Tests/extras.py
@@ -234,6 +234,20 @@ def custom_callback(name, progress):
 pointset_dict = itk.dict_from_pointset(pointset)
 pointset_back = itk.pointset_from_dict(pointset_dict)
 
+polyline = itk.PolyLineParametricPath[dim].New()
+polyline.AddVertex([0.0, 0.0])
+polyline.AddVertex([1.0, 1.0])
+polyline.AddVertex([4.0, 3.0])
+polyline_dict = itk.dict_from_polyline(polyline)
+polyline_back = itk.polyline_from_dict(polyline_dict)
+original_vertices = itk.array_from_vector_container(polyline.GetVertexList())
+back_vertices = itk.array_from_vector_container(polyline_back.GetVertexList())
+assert np.allclose(original_vertices, back_vertices)
+
+serialize_deserialize = pickle.loads(pickle.dumps(polyline))
+back_vertices = itk.array_from_vector_container(serialize_deserialize.GetVertexList())
+assert np.allclose(original_vertices, back_vertices)
+
 # test search
 res = itk.search("Index")
 assert res[0] == "Index"
diff --git a/Wrapping/Generators/Python/itk/support/extras.py b/Wrapping/Generators/Python/itk/support/extras.py
index bc58a0cf8e6..54343516e5b 100644
--- a/Wrapping/Generators/Python/itk/support/extras.py
+++ b/Wrapping/Generators/Python/itk/support/extras.py
@@ -103,6 +103,8 @@
     "dict_from_mesh",
     "pointset_from_dict",
     "dict_from_pointset",
+    "polyline_from_dict",
+    "dict_from_polyline",
     "transform_from_dict",
     "dict_from_transform",
     "transformwrite",
@@ -981,6 +983,30 @@ def dict_from_pointset(pointset: "itkt.PointSet") -> Dict:
     )
 
 
+def polyline_from_dict(polyline_dict: Dict) -> "itkt.PolylineParametricPath":
+    """Deserialize an dictionary representing an itk.PolylineParametricPath object."""
+    import itk
+
+    vertex_list = polyline_dict["vertexList"]
+    dimension = vertex_list.shape[1]
+    polyline = itk.PolyLineParametricPath[dimension].New()
+    polyline.SetObjectName(polyline_dict["name"])
+    for vertex in vertex_list:
+        polyline.AddVertex(vertex)
+
+    return polyline
+
+def dict_from_polyline(polyline: "itkt.PolylineParametricPath") -> Dict:
+    """Serialize a Python itk.PolylineParametricPath object to a pickable Python dictionary."""
+    import itk
+
+    vertex_list = polyline.GetVertexList()
+    vertex_list_array = itk.array_from_vector_container(vertex_list)
+    return dict(
+        name=polyline.GetObjectName(),
+        vertexList=vertex_list_array,
+    )
+
 def dict_from_transform(transform: "itkt.TransformBase") -> Dict:
     import itk
 
diff --git a/Wrapping/TypedefMacros.cmake b/Wrapping/TypedefMacros.cmake
index e04ae853e66..15dc1a8e2b9 100644
--- a/Wrapping/TypedefMacros.cmake
+++ b/Wrapping/TypedefMacros.cmake
@@ -639,6 +639,10 @@ macro(itk_wrap_simple_type wrap_class swig_name)
     string(APPEND ITK_WRAP_PYTHON_SWIG_EXT "DECL_PYTHON_POINTSET_CLASS(${swig_name})\n\n")
   endif()
 
+  if("${cpp_name}" STREQUAL "itk::PolyLineParametricPath")
+    string(APPEND ITK_WRAP_PYTHON_SWIG_EXT "DECL_PYTHON_POLYLINEPARAMETRICPATH_CLASS(${swig_name})\n\n")
+  endif()
+
   if("${cpp_name}" STREQUAL "itk::Mesh")
     string(APPEND ITK_WRAP_PYTHON_SWIG_EXT "DECL_PYTHON_MESH_CLASS(${swig_name})\n\n")
   endif()