diff --git a/raillabel/format/bbox.py b/raillabel/format/bbox.py index 682c704..b30b1e6 100644 --- a/raillabel/format/bbox.py +++ b/raillabel/format/bbox.py @@ -8,7 +8,7 @@ from raillabel.json_format import JSONBbox -from ._attributes import _attributes_from_json +from ._attributes import _attributes_from_json, _attributes_to_json from .point2d import Point2d from .size2d import Size2d @@ -43,6 +43,16 @@ def from_json(cls, json: JSONBbox, object_uid: UUID) -> Bbox: attributes=_attributes_from_json(json.attributes), ) + def to_json(self, uid: UUID, object_type: str) -> JSONBbox: + """Export this object into the RailLabel JSON format.""" + return JSONBbox( + name=self.name(object_type), + val=list(self.pos.to_json()) + list(self.size.to_json()), + coordinate_system=self.sensor, + 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}__bbox__{object_type}" diff --git a/tests/test_raillabel/format/test_bbox.py b/tests/test_raillabel/format/test_bbox.py index d6c670f..fb39bd2 100644 --- a/tests/test_raillabel/format/test_bbox.py +++ b/tests/test_raillabel/format/test_bbox.py @@ -62,5 +62,10 @@ def test_name(bbox): assert actual == "rgb_middle__bbox__person" +def test_to_json(bbox, bbox_json, bbox_uid): + actual = bbox.to_json(bbox_uid, object_type="person") + assert actual == bbox_json + + if __name__ == "__main__": pytest.main([__file__, "-v"])