feat: implement Poly2d.to_json()

raillabel/format/_util.py

@@ -0,0 +1,8 @@
+# Copyright DB InfraGO AG and contributors
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+
+def _flatten_list(list_of_tuples: list[tuple]) -> list:
+    return [item for tup in list_of_tuples for item in tup]

raillabel/format/poly2d.py

@@ -8,7 +8,8 @@
 
 from raillabel.json_format import JSONPoly2d
 
-from ._attributes import _attributes_from_json
+from ._attributes import _attributes_from_json, _attributes_to_json
+from ._util import _flatten_list
 from .point2d import Point2d
 
 
@@ -45,6 +46,18 @@ def from_json(cls, json: JSONPoly2d, object_id: UUID) -> Poly2d:
             attributes=_attributes_from_json(json.attributes),
         )
 
+    def to_json(self, uid: UUID, object_type: str) -> JSONPoly2d:
+        """Export this object into the RailLabel JSON format."""
+        return JSONPoly2d(
+            name=self.name(object_type),
+            val=_flatten_list([point.to_json() for point in self.points]),
+            closed=self.closed,
+            mode="MODE_POLY2D_ABSOLUTE",
+            coordinate_system=self.sensor_id,
+            uid=uid,
+            attributes=_attributes_to_json(self.attributes),
+        )
+
     def name(self, object_type: str) -> str:
         """Return the name of the annotation used for indexing in the object data pointers."""
         return f"{self.sensor_id}__poly2d__{object_type}"

tests/test_raillabel/format/test_poly2d.py

@@ -64,5 +64,10 @@ def test_name(poly2d):
     assert actual == "rgb_middle__poly2d__person"
 
 
+def test_to_json(poly2d, poly2d_json, poly2d_id):
+    actual = poly2d.to_json(poly2d_id, object_type="person")
+    assert actual == poly2d_json
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])