diff --git a/raillabel/format/seg3d.py b/raillabel/format/seg3d.py index 76ef176..31dd7c0 100644 --- a/raillabel/format/seg3d.py +++ b/raillabel/format/seg3d.py @@ -8,7 +8,7 @@ from raillabel.json_format import JSONVec -from ._attributes import _attributes_from_json +from ._attributes import _attributes_from_json, _attributes_to_json @dataclass @@ -37,6 +37,16 @@ def from_json(cls, json: JSONVec, object_id: UUID) -> Seg3d: attributes=_attributes_from_json(json.attributes), ) + def to_json(self, uid: UUID, object_type: str) -> JSONVec: + """Export this object into the RailLabel JSON format.""" + return JSONVec( + name=self.name(object_type), + val=self.point_ids, + 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}__vec__{object_type}" diff --git a/tests/test_raillabel/format/test_seg3d.py b/tests/test_raillabel/format/test_seg3d.py index 0dd2716..a290aa6 100644 --- a/tests/test_raillabel/format/test_seg3d.py +++ b/tests/test_raillabel/format/test_seg3d.py @@ -19,7 +19,7 @@ def seg3d_json( ) -> JSONVec: return JSONVec( uid="d52e2b25-0B48-4899-86d5-4bc41be6b7d3", - name="rgb_middle__seg3d__person", + name="lidar__vec__person", val=[1234, 5678], coordinate_system="lidar", attributes=attributes_multiple_types_json, @@ -57,5 +57,10 @@ def test_name(seg3d): assert actual == "lidar__vec__person" +def test_to_json(seg3d, seg3d_json, seg3d_id): + actual = seg3d.to_json(seg3d_id, object_type="person") + assert actual == seg3d_json + + if __name__ == "__main__": pytest.main([__file__, "-v"])