Add confidence to SingleInstanceSegmentation

src/labelformat/formats/coco.py

@@ -134,6 +134,7 @@ def get_labels(self) -> Iterable[ImageInstanceSegmentation]:
                     SingleInstanceSegmentation(
                         category=category_id_to_category[ann["category_id"]],
                         segmentation=segmentation,
+                        confidence=(float(ann["score"]) if "score" in ann else None),
                     )
                 )
             yield ImageInstanceSegmentation(
@@ -204,6 +205,8 @@ def save(self, label_input: InstanceSegmentationInput) -> None:
                     "iscrowd": is_crowd,
                     "segmentation": segmentation,
                 }
+                if obj.confidence is not None:
+                    annotation["score"] = obj.confidence
                 data["annotations"].append(annotation)
 
         self.output_file.parent.mkdir(parents=True, exist_ok=True)

src/labelformat/model/instance_segmentation.py

@@ -15,6 +15,13 @@
 class SingleInstanceSegmentation:
     category: Category
     segmentation: MultiPolygon | BinaryMaskSegmentation
+    confidence: float | None = None
+
+    def __post_init__(self) -> None:
+        if self.confidence is not None and not (0 <= self.confidence <= 1):
+            raise ValueError(
+                f"Confidence must be between 0 and 1, but got: {self.confidence}"
+            )
 
 
 @dataclass(frozen=True)

tests/unit/formats/test_coco.py

@@ -3,7 +3,12 @@
 
 import pytest
 
-from labelformat.formats.coco import COCOObjectDetectionInput, COCOObjectDetectionOutput
+from labelformat.formats.coco import (
+    COCOInstanceSegmentationInput,
+    COCOInstanceSegmentationOutput,
+    COCOObjectDetectionInput,
+    COCOObjectDetectionOutput,
+)
 from labelformat.model.bounding_box import BoundingBox
 from labelformat.model.category import Category
 from labelformat.model.image import Image
@@ -99,3 +104,40 @@ def test_save(self, tmp_path: Path, with_confidence: bool) -> None:
             expected_annotations[0]["score"] = 0.4
             expected_annotations[1]["score"] = 0.8
         assert output_json["annotations"] == expected_annotations
+
+
+class TestCOCOInstanceSegmentationConfidence:
+    @pytest.mark.parametrize("with_score", [True, False])
+    def test_round_trip(self, tmp_path: Path, with_score: bool) -> None:
+        annotation = {
+            "image_id": 0,
+            "category_id": 1,
+            "bbox": [10.0, 10.0, 10.0, 10.0],
+            "iscrowd": 0,
+            "segmentation": [[10.0, 10.0, 10.0, 20.0, 20.0, 20.0, 20.0, 10.0]],
+        }
+        if with_score:
+            annotation["score"] = 0.4
+        input_file = tmp_path / "input.json"
+        input_file.write_text(
+            json.dumps(
+                {
+                    "images": [
+                        {"id": 0, "file_name": "image.jpg", "width": 100, "height": 200}
+                    ],
+                    "categories": [{"id": 1, "name": "dog"}],
+                    "annotations": [annotation],
+                }
+            )
+        )
+
+        label_input = COCOInstanceSegmentationInput(input_file=input_file)
+        [label] = label_input.get_labels()
+        assert label.objects[0].confidence == (0.4 if with_score else None)
+
+        output_file = tmp_path / "output.json"
+        COCOInstanceSegmentationOutput(output_file=output_file).save(
+            label_input=label_input
+        )
+        [saved] = json.loads(output_file.read_text())["annotations"]
+        assert saved.get("score") == annotation.get("score")

tests/unit/model/test_instance_segmentation.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import pytest
+
+from labelformat.model.category import Category
+from labelformat.model.instance_segmentation import SingleInstanceSegmentation
+from labelformat.model.multipolygon import MultiPolygon
+
+_SQUARE = MultiPolygon(polygons=[[(0.0, 0.0), (0.0, 1.0), (1.0, 1.0), (1.0, 0.0)]])
+
+
+class TestSingleInstanceSegmentation:
+    @pytest.mark.parametrize("confidence", [None, 0.0, 1.0])
+    def test_confidence_valid(self, confidence: float | None) -> None:
+        instance = SingleInstanceSegmentation(
+            category=Category(id=0, name="cat"),
+            segmentation=_SQUARE,
+            confidence=confidence,
+        )
+        assert instance.confidence == confidence
+
+    @pytest.mark.parametrize(
+        "confidence, expected_error",
+        [
+            (-0.1, "Confidence must be between 0 and 1, but got: -0.1"),
+            (1.1, "Confidence must be between 0 and 1, but got: 1.1"),
+        ],
+    )
+    def test_confidence_out_of_bounds(
+        self, confidence: float, expected_error: str
+    ) -> None:
+        with pytest.raises(ValueError, match=expected_error):
+            SingleInstanceSegmentation(
+                category=Category(id=0, name="cat"),
+                segmentation=_SQUARE,
+                confidence=confidence,
+            )