Merge branch 'jonas-lig-8150-instance-segmentation-youtubevis-format-coco-helpers' of github.com:lightly-ai/labelformat into jonas-lig-8150-instance-segmentation-youtubevis-format-coco-helpers

Author: JonasWurst

src/labelformat/model/instance_segmentation_track.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from argparse import ArgumentParser
+from dataclasses import dataclass
+from typing import Iterable
+
+from labelformat.model.binary_mask_segmentation import BinaryMaskSegmentation
+from labelformat.model.category import Category
+from labelformat.model.multipolygon import MultiPolygon
+from labelformat.model.video import Video
+
+
+@dataclass(frozen=True)
+class SingleInstanceSegmentationTrack:
+    category: Category
+    segmentations: list[MultiPolygon | BinaryMaskSegmentation | None]
+
+
+@dataclass(frozen=True)
+class VideoInstanceSegmentationTrack:
+    """
+    The base class for a video alongside with its object detection track annotations.
+    A video consists of N frames and M objects. Each object is defined by N boxes - one for each frame.
+    If an object is not present on a frame, the corresponding entry is set to None.
+    """
+
+    video: Video
+    objects: list[SingleInstanceSegmentationTrack]
+
+    def __post_init__(self) -> None:
+        number_of_frames = self.video.number_of_frames
+
+        for obj in self.objects:
+            if len(obj.segmentations) != number_of_frames:
+                raise ValueError(
+                    "Length of instance segmentation track does not match the number of frames in the video."
+                )
+
+
+class InstanceSegmentationTrackInput(ABC):
+    @staticmethod
+    @abstractmethod
+    def add_cli_arguments(parser: ArgumentParser) -> None:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_categories(self) -> Iterable[Category]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_videos(self) -> Iterable[Video]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_labels(self) -> Iterable[VideoInstanceSegmentationTrack]:
+        raise NotImplementedError()
+
+
+class InstanceSegmentationTrackOutput(ABC):
+    @staticmethod
+    @abstractmethod
+    def add_cli_arguments(parser: ArgumentParser) -> None:
+        raise NotImplementedError()
+
+    def save(self, label_input: InstanceSegmentationTrackInput) -> None:
+        raise NotImplementedError()

tests/unit/model/test_instance_segmentation_track.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import pytest
+
+from labelformat.model.category import Category
+from labelformat.model.instance_segmentation_track import (
+    SingleInstanceSegmentationTrack,
+    VideoInstanceSegmentationTrack,
+)
+from labelformat.model.multipolygon import MultiPolygon
+from labelformat.model.video import Video
+
+
+class TestVideoInstanceSegmentationTrack:
+    def test_post_init__frames_equal_segmentations_length__valid(self) -> None:
+        track_a = SingleInstanceSegmentationTrack(
+            category=Category(id=0, name="cat"),
+            segmentations=[
+                MultiPolygon(polygons=[[(0.0, 0.0), (1.0, 0.0), (1.0, 1.0)]]),
+                None,
+            ],
+        )
+
+        track_b = SingleInstanceSegmentationTrack(
+            category=Category(id=1, name="dog"),
+            segmentations=[
+                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)]]),
+            ],
+        )
+
+        video = Video(id=0, filename="test.mov", width=1, height=1, number_of_frames=2)
+
+        instance_seg = VideoInstanceSegmentationTrack(
+            video=video,
+            objects=[track_a, track_b],
+        )
+        assert len(instance_seg.objects) == 2
+        assert len(instance_seg.objects[0].segmentations) == 2
+
+    def test_post_init__frames_equal_segmentations_length___invalid(self) -> None:
+        track_a = SingleInstanceSegmentationTrack(
+            category=Category(id=0, name="cat"),
+            segmentations=[
+                MultiPolygon(polygons=[[(0.0, 0.0), (1.0, 0.0), (1.0, 1.0)]]),
+                None,
+                None,
+            ],
+        )
+
+        video = Video(id=0, filename="test.mov", width=1, height=1, number_of_frames=2)
+
+        with pytest.raises(
+            ValueError,
+            match="Length of instance segmentation track does not match the number of frames in the video.",
+        ):
+            VideoInstanceSegmentationTrack(
+                video=video,
+                objects=[track_a],
+            )