Add object track id to youtubevis annotation models (#68)

Author: horatiualmasan

src/labelformat/formats/youtubevis.py

@@ -84,6 +84,7 @@ def get_labels(self) -> Iterable[VideoObjectDetectionTrack]:
                     SingleObjectDetectionTrack(
                         category=category_id_to_category[track["category_id"]],
                         boxes=boxes,
+                        object_track_id=track.get("id"),
                     )
                 )
             yield VideoObjectDetectionTrack(
@@ -115,6 +116,7 @@ def get_labels(self) -> Iterable[VideoInstanceSegmentationTrack]:
                     SingleInstanceSegmentationTrack(
                         category=category_id_to_category[track["category_id"]],
                         segmentations=segmentations,
+                        object_track_id=track.get("id"),
                     )
                 )
             yield VideoInstanceSegmentationTrack(

src/labelformat/model/instance_segmentation_track.py

@@ -15,6 +15,7 @@
 class SingleInstanceSegmentationTrack:
     category: Category
     segmentations: list[MultiPolygon | BinaryMaskSegmentation | None]
+    object_track_id: int | None = None
 
 
 @dataclass(frozen=True)

src/labelformat/model/object_detection_track.py

@@ -14,6 +14,7 @@
 class SingleObjectDetectionTrack:
     category: Category
     boxes: list[BoundingBox | None]
+    object_track_id: int | None = None
     # TODO (Jonas, 01/2026): Add confidence
 
 

tests/unit/formats/test_youtubevis.py

@@ -70,6 +70,7 @@ def test_get_labels(self, tmp_path: Path) -> None:
                             ),
                             None,
                         ],
+                        object_track_id=10,
                     )
                 ],
             )
@@ -103,6 +104,7 @@ def test_get_labels(self, tmp_path: Path) -> None:
                     SingleInstanceSegmentationTrack(
                         category=Category(id=1, name="cat"),
                         segmentations=[expected_segmentation, None],
+                        object_track_id=20,
                     )
                 ],
             )
@@ -125,6 +127,7 @@ def _write_youtube_vis_json(input_file: Path) -> Path:
         ],
         "annotations": [
             {
+                "id": 10,
                 "video_id": 5,
                 "category_id": 1,
                 "bboxes": [
@@ -156,6 +159,7 @@ def _write_youtube_vis_instance_segmentation_json(input_file: Path) -> Path:
         ],
         "annotations": [
             {
+                "id": 20,
                 "video_id": 5,
                 "category_id": 1,
                 "bboxes": [

tests/unit/model/test_instance_segmentation_track.py

@@ -19,6 +19,7 @@ def test_post_init__frames_equal_segmentations_length__valid(self) -> None:
                 MultiPolygon(polygons=[[(0.0, 0.0), (1.0, 0.0), (1.0, 1.0)]]),
                 None,
             ],
+            object_track_id=0,
         )
 
         track_b = SingleInstanceSegmentationTrack(
@@ -27,6 +28,7 @@ def test_post_init__frames_equal_segmentations_length__valid(self) -> None:
                 MultiPolygon(polygons=[[(2.0, 2.0), (3.0, 2.0), (3.0, 3.0)]]),
                 MultiPolygon(polygons=[[(4.0, 4.0), (5.0, 4.0), (5.0, 5.0)]]),
             ],
+            object_track_id=1,
         )
 
         video = Video(id=0, filename="test.mov", width=1, height=1, number_of_frames=2)
@@ -46,6 +48,7 @@ def test_post_init__frames_equal_segmentations_length___invalid(self) -> None:
                 None,
                 None,
             ],
+            object_track_id=0,
         )
 
         video = Video(id=0, filename="test.mov", width=1, height=1, number_of_frames=2)

tests/unit/model/test_object_detection_track.py

@@ -16,11 +16,13 @@ def test_post_init__frames_equal_boxes_length__valid(self) -> None:
         track_a = SingleObjectDetectionTrack(
             category=Category(id=0, name="cat"),
             boxes=[BoundingBox(xmin=0, ymin=0, xmax=1, ymax=1) for _ in range(2)],
+            object_track_id=0,
         )
 
         track_b = SingleObjectDetectionTrack(
             category=Category(id=1, name="dog"),
             boxes=[BoundingBox(xmin=0, ymin=0, xmax=1, ymax=1) for _ in range(2)],
+            object_track_id=1,
         )
 
         video = Video(id=0, filename="test.mov", width=1, height=1, number_of_frames=2)
@@ -36,11 +38,13 @@ def test_post_init__frames_equal_boxes_length___invalid(self) -> None:
         track_a = SingleObjectDetectionTrack(
             category=Category(id=0, name="cat"),
             boxes=[BoundingBox(xmin=0, ymin=0, xmax=1, ymax=1) for _ in range(2)],
+            object_track_id=0,
         )
 
         track_b = SingleObjectDetectionTrack(
             category=Category(id=1, name="dog"),
             boxes=[BoundingBox(xmin=0, ymin=0, xmax=1, ymax=1) for _ in range(3)],
+            object_track_id=1,
         )
 
         video = Video(id=0, filename="test.mov", width=1, height=1, number_of_frames=2)