From d1bd3b977b6258ad5bca95ca5f1889a9b42682b5 Mon Sep 17 00:00:00 2001
From: unexcellent <>
Date: Thu, 14 Nov 2024 06:08:25 +0100
Subject: [PATCH] fix: extra fields allowed in json format classes

---
 raillabel/format/frame.py                     | 10 ++++-----
 raillabel/json_format/__init__.py             |  3 ++-
 raillabel/json_format/attributes.py           |  2 +-
 raillabel/json_format/bbox.py                 |  2 +-
 raillabel/json_format/boolean_attribute.py    |  2 +-
 raillabel/json_format/coordinate_system.py    |  2 +-
 raillabel/json_format/cuboid.py               |  2 +-
 raillabel/json_format/element_data_pointer.py |  2 +-
 raillabel/json_format/frame.py                |  6 ++---
 raillabel/json_format/frame_interval.py       |  2 +-
 raillabel/json_format/num.py                  |  2 +-
 raillabel/json_format/num_attribute.py        |  2 +-
 raillabel/json_format/object.py               |  2 +-
 raillabel/json_format/object_data.py          |  8 ++++++-
 raillabel/json_format/poly2d.py               |  2 +-
 raillabel/json_format/poly3d.py               |  2 +-
 raillabel/json_format/scene.py                |  4 ++--
 raillabel/json_format/stream_camera.py        |  6 ++---
 raillabel/json_format/stream_other.py         |  2 +-
 raillabel/json_format/stream_radar.py         |  6 ++---
 raillabel/json_format/stream_sync.py          |  6 ++---
 raillabel/json_format/text_attribute.py       |  2 +-
 raillabel/json_format/transform_data.py       |  2 +-
 raillabel/json_format/vec.py                  |  2 +-
 raillabel/json_format/vec_attribute.py        |  2 +-
 tests/test_raillabel/format/conftest.py       |  2 +-
 tests/test_raillabel/format/test_frame.py     | 22 ++++++++++++++-----
 27 files changed, 62 insertions(+), 45 deletions(-)

diff --git a/raillabel/format/frame.py b/raillabel/format/frame.py
index 182523b..686715e 100644
--- a/raillabel/format/frame.py
+++ b/raillabel/format/frame.py
@@ -91,19 +91,19 @@ def _annotations_from_json(
     annotations: dict[UUID, Bbox | Cuboid | Poly2d | Poly3d | Seg3d] = {}
 
     for object_id, object_data in json_object_data.items():
-        for json_bbox in _resolve_none_to_empty_list(object_data.bbox):
+        for json_bbox in _resolve_none_to_empty_list(object_data.object_data.bbox):
             annotations[json_bbox.uid] = Bbox.from_json(json_bbox, object_id)
 
-        for json_cuboid in _resolve_none_to_empty_list(object_data.cuboid):
+        for json_cuboid in _resolve_none_to_empty_list(object_data.object_data.cuboid):
             annotations[json_cuboid.uid] = Cuboid.from_json(json_cuboid, object_id)
 
-        for json_poly2d in _resolve_none_to_empty_list(object_data.poly2d):
+        for json_poly2d in _resolve_none_to_empty_list(object_data.object_data.poly2d):
             annotations[json_poly2d.uid] = Poly2d.from_json(json_poly2d, object_id)
 
-        for json_poly3d in _resolve_none_to_empty_list(object_data.poly3d):
+        for json_poly3d in _resolve_none_to_empty_list(object_data.object_data.poly3d):
             annotations[json_poly3d.uid] = Poly3d.from_json(json_poly3d, object_id)
 
-        for json_seg3d in _resolve_none_to_empty_list(object_data.vec):
+        for json_seg3d in _resolve_none_to_empty_list(object_data.object_data.vec):
             annotations[json_seg3d.uid] = Seg3d.from_json(json_seg3d, object_id)
 
     return annotations
diff --git a/raillabel/json_format/__init__.py b/raillabel/json_format/__init__.py
index 059d7d1..9ad9f51 100644
--- a/raillabel/json_format/__init__.py
+++ b/raillabel/json_format/__init__.py
@@ -14,7 +14,7 @@
 from .num import JSONNum
 from .num_attribute import JSONNumAttribute
 from .object import JSONObject
-from .object_data import JSONObjectData
+from .object_data import JSONAnnotations, JSONObjectData
 from .poly2d import JSONPoly2d
 from .poly3d import JSONPoly3d
 from .scene import JSONScene, JSONSceneContent
@@ -28,6 +28,7 @@
 from .vec_attribute import JSONVecAttribute
 
 __all__ = [
+    "JSONAnnotations",
     "JSONAttributes",
     "JSONBbox",
     "JSONBooleanAttribute",
diff --git a/raillabel/json_format/attributes.py b/raillabel/json_format/attributes.py
index 976c44d..68aaacb 100644
--- a/raillabel/json_format/attributes.py
+++ b/raillabel/json_format/attributes.py
@@ -11,7 +11,7 @@
 from .vec_attribute import JSONVecAttribute
 
 
-class JSONAttributes(BaseModel):
+class JSONAttributes(BaseModel, extra="forbid"):
     """Attributes is the alias of element data that can be nested inside geometric object data.
 
     For example, a certain bounding box can have attributes related to its score, visibility, etc.
diff --git a/raillabel/json_format/bbox.py b/raillabel/json_format/bbox.py
index f82c8c5..944db2c 100644
--- a/raillabel/json_format/bbox.py
+++ b/raillabel/json_format/bbox.py
@@ -10,7 +10,7 @@
 from .attributes import JSONAttributes
 
 
-class JSONBbox(BaseModel):
+class JSONBbox(BaseModel, extra="forbid"):
     """A 2D bounding box is defined as a 4-dimensional vector [x, y, w, h].
 
     [x, y] is the center of the bounding box and [w, h] represent the width (horizontal,
diff --git a/raillabel/json_format/boolean_attribute.py b/raillabel/json_format/boolean_attribute.py
index afda14b..91c86ea 100644
--- a/raillabel/json_format/boolean_attribute.py
+++ b/raillabel/json_format/boolean_attribute.py
@@ -6,7 +6,7 @@
 from pydantic import BaseModel
 
 
-class JSONBooleanAttribute(BaseModel):
+class JSONBooleanAttribute(BaseModel, extra="forbid"):
     """A boolean attribute."""
 
     name: str
diff --git a/raillabel/json_format/coordinate_system.py b/raillabel/json_format/coordinate_system.py
index e1d664a..fedc174 100644
--- a/raillabel/json_format/coordinate_system.py
+++ b/raillabel/json_format/coordinate_system.py
@@ -10,7 +10,7 @@
 from .transform_data import JSONTransformData
 
 
-class JSONCoordinateSystem(BaseModel):
+class JSONCoordinateSystem(BaseModel, extra="forbid"):
     """A 3D reference frame."""
 
     parent: Literal["base", ""]
diff --git a/raillabel/json_format/cuboid.py b/raillabel/json_format/cuboid.py
index 1a918a4..a58431a 100644
--- a/raillabel/json_format/cuboid.py
+++ b/raillabel/json_format/cuboid.py
@@ -10,7 +10,7 @@
 from .attributes import JSONAttributes
 
 
-class JSONCuboid(BaseModel):
+class JSONCuboid(BaseModel, extra="forbid"):
     """A cuboid or 3D bounding box.
 
     It is defined by the position of its center, the rotation in 3D, and its dimensions.
diff --git a/raillabel/json_format/element_data_pointer.py b/raillabel/json_format/element_data_pointer.py
index 48cfa14..745c9c0 100644
--- a/raillabel/json_format/element_data_pointer.py
+++ b/raillabel/json_format/element_data_pointer.py
@@ -10,7 +10,7 @@
 from .frame_interval import JSONFrameInterval
 
 
-class JSONElementDataPointer(BaseModel):
+class JSONElementDataPointer(BaseModel, extra="forbid"):
     """A pointer to element data of elements.
 
     It is indexed by 'name', and containing information about the element data type, for example,
diff --git a/raillabel/json_format/frame.py b/raillabel/json_format/frame.py
index 797baec..cfc1386 100644
--- a/raillabel/json_format/frame.py
+++ b/raillabel/json_format/frame.py
@@ -13,7 +13,7 @@
 from .stream_sync import JSONStreamSync
 
 
-class JSONFrame(BaseModel):
+class JSONFrame(BaseModel, extra="forbid"):
     """A frame is a container of dynamic, timewise, information."""
 
     frame_properties: JSONFrameProperties | None = None
@@ -24,7 +24,7 @@ class JSONFrame(BaseModel):
     strings containing 32 bytes UUIDs. Object values contain an 'object_data' JSON object."""
 
 
-class JSONFrameProperties(BaseModel):
+class JSONFrameProperties(BaseModel, extra="forbid"):
     """Container of frame information other than annotations."""
 
     timestamp: Decimal | str | None = None
@@ -38,7 +38,7 @@ class JSONFrameProperties(BaseModel):
     "Additional data to describe attributes of the frame (like GPS position)."
 
 
-class JSONFrameData(BaseModel):
+class JSONFrameData(BaseModel, extra="forbid"):
     """Additional data to describe attributes of the frame (like GPS position)."""
 
     num: list[JSONNum] | None = None
diff --git a/raillabel/json_format/frame_interval.py b/raillabel/json_format/frame_interval.py
index b56cbe9..d7f9034 100644
--- a/raillabel/json_format/frame_interval.py
+++ b/raillabel/json_format/frame_interval.py
@@ -6,7 +6,7 @@
 from pydantic import BaseModel
 
 
-class JSONFrameInterval(BaseModel):
+class JSONFrameInterval(BaseModel, extra="forbid"):
     """A frame interval defines a starting and ending frame number as a closed interval.
 
     That means the interval includes the limit frame numbers.
diff --git a/raillabel/json_format/num.py b/raillabel/json_format/num.py
index b51ac76..be9089b 100644
--- a/raillabel/json_format/num.py
+++ b/raillabel/json_format/num.py
@@ -8,7 +8,7 @@
 from pydantic import BaseModel
 
 
-class JSONNum(BaseModel):
+class JSONNum(BaseModel, extra="forbid"):
     """A number."""
 
     name: str
diff --git a/raillabel/json_format/num_attribute.py b/raillabel/json_format/num_attribute.py
index 700209c..e829712 100644
--- a/raillabel/json_format/num_attribute.py
+++ b/raillabel/json_format/num_attribute.py
@@ -6,7 +6,7 @@
 from pydantic import BaseModel
 
 
-class JSONNumAttribute(BaseModel):
+class JSONNumAttribute(BaseModel, extra="forbid"):
     """A number attribute."""
 
     name: str
diff --git a/raillabel/json_format/object.py b/raillabel/json_format/object.py
index 091843a..514f1e8 100644
--- a/raillabel/json_format/object.py
+++ b/raillabel/json_format/object.py
@@ -9,7 +9,7 @@
 from .frame_interval import JSONFrameInterval
 
 
-class JSONObject(BaseModel):
+class JSONObject(BaseModel, extra="forbid"):
     """An object is the main type of annotation element.
 
     Object is designed to represent spatiotemporal entities, such as physical objects in the real
diff --git a/raillabel/json_format/object_data.py b/raillabel/json_format/object_data.py
index bffa98a..04609e7 100644
--- a/raillabel/json_format/object_data.py
+++ b/raillabel/json_format/object_data.py
@@ -12,9 +12,15 @@
 from .vec import JSONVec
 
 
-class JSONObjectData(BaseModel):
+class JSONObjectData(BaseModel, extra="forbid"):
     """Container of annotations of an object in a frame."""
 
+    object_data: JSONAnnotations
+
+
+class JSONAnnotations(BaseModel, extra="forbid"):
+    """Container of the annotations by type."""
+
     bbox: list[JSONBbox] | None = None
     cuboid: list[JSONCuboid] | None = None
     poly2d: list[JSONPoly2d] | None = None
diff --git a/raillabel/json_format/poly2d.py b/raillabel/json_format/poly2d.py
index 9d31a16..5ccda8c 100644
--- a/raillabel/json_format/poly2d.py
+++ b/raillabel/json_format/poly2d.py
@@ -11,7 +11,7 @@
 from .attributes import JSONAttributes
 
 
-class JSONPoly2d(BaseModel):
+class JSONPoly2d(BaseModel, extra="forbid"):
     """A 2D polyline defined as a sequence of 2D points."""
 
     name: str
diff --git a/raillabel/json_format/poly3d.py b/raillabel/json_format/poly3d.py
index eb72dd2..714204a 100644
--- a/raillabel/json_format/poly3d.py
+++ b/raillabel/json_format/poly3d.py
@@ -10,7 +10,7 @@
 from .attributes import JSONAttributes
 
 
-class JSONPoly3d(BaseModel):
+class JSONPoly3d(BaseModel, extra="forbid"):
     """A 3D polyline defined as a sequence of 3D points."""
 
     name: str
diff --git a/raillabel/json_format/scene.py b/raillabel/json_format/scene.py
index 65326ec..36135d6 100644
--- a/raillabel/json_format/scene.py
+++ b/raillabel/json_format/scene.py
@@ -17,13 +17,13 @@
 from .stream_radar import JSONStreamRadar
 
 
-class JSONScene(BaseModel):
+class JSONScene(BaseModel, extra="forbid"):
     """Root RailLabel object."""
 
     openlabel: JSONSceneContent
 
 
-class JSONSceneContent(BaseModel):
+class JSONSceneContent(BaseModel, extra="forbid"):
     """Container of all scene content."""
 
     metadata: JSONMetadata
diff --git a/raillabel/json_format/stream_camera.py b/raillabel/json_format/stream_camera.py
index b8f27ec..1541c64 100644
--- a/raillabel/json_format/stream_camera.py
+++ b/raillabel/json_format/stream_camera.py
@@ -8,7 +8,7 @@
 from pydantic import BaseModel
 
 
-class JSONStreamCamera(BaseModel):
+class JSONStreamCamera(BaseModel, extra="forbid"):
     """A stream describes the source of a data sequence, usually a sensor.
 
     This specific object contains the intrinsics of a camera sensor.
@@ -27,13 +27,13 @@ class JSONStreamCamera(BaseModel):
     "Description of the stream."
 
 
-class JSONStreamCameraProperties(BaseModel):
+class JSONStreamCameraProperties(BaseModel, extra="forbid"):
     """Intrinsic calibration of the stream."""
 
     intrinsics_pinhole: JSONIntrinsicsPinhole
 
 
-class JSONIntrinsicsPinhole(BaseModel):
+class JSONIntrinsicsPinhole(BaseModel, extra="forbid"):
     """JSON object defining an instance of the intrinsic parameters of a pinhole camera."""
 
     camera_matrix: tuple[
diff --git a/raillabel/json_format/stream_other.py b/raillabel/json_format/stream_other.py
index a1d47ba..75fd241 100644
--- a/raillabel/json_format/stream_other.py
+++ b/raillabel/json_format/stream_other.py
@@ -8,7 +8,7 @@
 from pydantic import BaseModel
 
 
-class JSONStreamOther(BaseModel):
+class JSONStreamOther(BaseModel, extra="forbid"):
     """A stream describes the source of a data sequence, usually a sensor.
 
     This specific object describes a sensor without intrinsic calibration.
diff --git a/raillabel/json_format/stream_radar.py b/raillabel/json_format/stream_radar.py
index 3c4a19c..1ac48e3 100644
--- a/raillabel/json_format/stream_radar.py
+++ b/raillabel/json_format/stream_radar.py
@@ -8,7 +8,7 @@
 from pydantic import BaseModel
 
 
-class JSONStreamRadar(BaseModel):
+class JSONStreamRadar(BaseModel, extra="forbid"):
     """A stream describes the source of a data sequence, usually a sensor.
 
     This specific object contains the intrinsics of a radar sensor.
@@ -27,13 +27,13 @@ class JSONStreamRadar(BaseModel):
     "Description of the stream."
 
 
-class JSONStreamRadarProperties(BaseModel):
+class JSONStreamRadarProperties(BaseModel, extra="forbid"):
     """Intrinsic calibration of the stream."""
 
     intrinsics_radar: JSONIntrinsicsRadar
 
 
-class JSONIntrinsicsRadar(BaseModel):
+class JSONIntrinsicsRadar(BaseModel, extra="forbid"):
     """JSON object defining an instance of the intrinsic parameters of a radar."""
 
     resolution_px_per_m: float
diff --git a/raillabel/json_format/stream_sync.py b/raillabel/json_format/stream_sync.py
index 2757924..5e48974 100644
--- a/raillabel/json_format/stream_sync.py
+++ b/raillabel/json_format/stream_sync.py
@@ -8,20 +8,20 @@
 from pydantic import BaseModel
 
 
-class JSONStreamSync(BaseModel):
+class JSONStreamSync(BaseModel, extra="forbid"):
     """Syncronization information of a stream in a frame."""
 
     stream_properties: JSONStreamSyncProperties
     uri: str | None = None
 
 
-class JSONStreamSyncProperties(BaseModel):
+class JSONStreamSyncProperties(BaseModel, extra="forbid"):
     """The sync information."""
 
     sync: JSONStreamSyncTimestamp
 
 
-class JSONStreamSyncTimestamp(BaseModel):
+class JSONStreamSyncTimestamp(BaseModel, extra="forbid"):
     """The timestamp of a stream sync."""
 
     timestamp: Decimal | str
diff --git a/raillabel/json_format/text_attribute.py b/raillabel/json_format/text_attribute.py
index cb1fb23..24cc389 100644
--- a/raillabel/json_format/text_attribute.py
+++ b/raillabel/json_format/text_attribute.py
@@ -6,7 +6,7 @@
 from pydantic import BaseModel
 
 
-class JSONTextAttribute(BaseModel):
+class JSONTextAttribute(BaseModel, extra="forbid"):
     """A text attribute."""
 
     name: str
diff --git a/raillabel/json_format/transform_data.py b/raillabel/json_format/transform_data.py
index edf7c87..4d5a235 100644
--- a/raillabel/json_format/transform_data.py
+++ b/raillabel/json_format/transform_data.py
@@ -6,7 +6,7 @@
 from pydantic import BaseModel
 
 
-class JSONTransformData(BaseModel):
+class JSONTransformData(BaseModel, extra="forbid"):
     """The translation and rotation of one coordinate system to another."""
 
     translation: tuple[float, float, float]
diff --git a/raillabel/json_format/vec.py b/raillabel/json_format/vec.py
index daa01c8..d9dec89 100644
--- a/raillabel/json_format/vec.py
+++ b/raillabel/json_format/vec.py
@@ -11,7 +11,7 @@
 from .attributes import JSONAttributes
 
 
-class JSONVec(BaseModel):
+class JSONVec(BaseModel, extra="forbid"):
     """A vector (list) of numbers."""
 
     name: str
diff --git a/raillabel/json_format/vec_attribute.py b/raillabel/json_format/vec_attribute.py
index 49abce7..c6df6a2 100644
--- a/raillabel/json_format/vec_attribute.py
+++ b/raillabel/json_format/vec_attribute.py
@@ -6,7 +6,7 @@
 from pydantic import BaseModel
 
 
-class JSONVecAttribute(BaseModel):
+class JSONVecAttribute(BaseModel, extra="forbid"):
     """A vec attribute."""
 
     name: str
diff --git a/tests/test_raillabel/format/conftest.py b/tests/test_raillabel/format/conftest.py
index 0bc295a..8ba0856 100644
--- a/tests/test_raillabel/format/conftest.py
+++ b/tests/test_raillabel/format/conftest.py
@@ -10,7 +10,7 @@
 from .test_intrinsics_radar import intrinsics_radar, intrinsics_radar_json
 from .test_lidar import lidar, lidar_json
 from .test_metadata import metadata, metadata_json
-from .test_num import num, num_json
+from .test_num import num, num_json, num_id
 from .test_object import (
     object_person,
     object_person_json,
diff --git a/tests/test_raillabel/format/test_frame.py b/tests/test_raillabel/format/test_frame.py
index 44814c1..07cd720 100644
--- a/tests/test_raillabel/format/test_frame.py
+++ b/tests/test_raillabel/format/test_frame.py
@@ -8,7 +8,13 @@
 import pytest
 
 from raillabel.format import Frame
-from raillabel.json_format import JSONFrame, JSONFrameData, JSONFrameProperties, JSONObjectData
+from raillabel.json_format import (
+    JSONFrame,
+    JSONFrameData,
+    JSONFrameProperties,
+    JSONObjectData,
+    JSONAnnotations,
+)
 
 # == Fixtures =========================
 
@@ -35,13 +41,17 @@ def frame_json(
         ),
         objects={
             "cfcf9750-3bc3-4077-9079-a82c0c63976a": JSONObjectData(
-                poly2d=[poly2d_json],
-                poly3d=[poly3d_json],
+                object_data=JSONAnnotations(
+                    poly2d=[poly2d_json],
+                    poly3d=[poly3d_json],
+                )
             ),
             "b40ba3ad-0327-46ff-9c28-2506cfd6d934": JSONObjectData(
-                bbox=[bbox_json],
-                cuboid=[cuboid_json],
-                vec=[seg3d_json],
+                object_data=JSONAnnotations(
+                    bbox=[bbox_json],
+                    cuboid=[cuboid_json],
+                    vec=[seg3d_json],
+                )
             ),
         },
     )