Skip to content

Commit

Permalink
feat: implement Poly3d.to_json()
Browse files Browse the repository at this point in the history
  • Loading branch information
unexcellent committed Nov 14, 2024
1 parent 342f627 commit 6d354b5
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 1 deletion.
14 changes: 13 additions & 1 deletion raillabel/format/poly3d.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,8 @@

from raillabel.json_format import JSONPoly3d

from ._attributes import _attributes_from_json
from ._attributes import _attributes_from_json, _attributes_to_json
from ._util import _flatten_list
from .point3d import Point3d


Expand Down Expand Up @@ -45,6 +46,17 @@ def from_json(cls, json: JSONPoly3d, object_id: UUID) -> Poly3d:
attributes=_attributes_from_json(json.attributes),
)

def to_json(self, uid: UUID, object_type: str) -> JSONPoly3d:
"""Export this object into the RailLabel JSON format."""
return JSONPoly3d(
name=self.name(object_type),
val=_flatten_list([point.to_json() for point in self.points]),
closed=self.closed,
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}__poly3d__{object_type}"
5 changes: 5 additions & 0 deletions tests/test_raillabel/format/test_poly3d.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,5 +63,10 @@ def test_name(poly3d):
assert actual == "lidar__poly3d__person"


def test_to_json(poly3d, poly3d_json, poly3d_id):
actual = poly3d.to_json(poly3d_id, object_type="person")
assert actual == poly3d_json


if __name__ == "__main__":
pytest.main([__file__, "-v"])

0 comments on commit 6d354b5

Please sign in to comment.