Lite Semantic Segmentation (#773)

Co-authored-by: Nick L <[email protected]>

.github/workflows/lite-tests-and-coverage.yml

@@ -36,7 +36,21 @@ jobs:
       - name: run object detection tests and report coverage
         run: |
           pip install -e ".[test]"
-          COVERAGE_FILE=.coverage.functional python -m coverage run --omit "tests/*"  -m pytest -v tests/detection/
+          COVERAGE_FILE=.coverage.detection python -m coverage run --omit "tests/*"  -m pytest -v tests/detection/
+          python -m coverage combine
+          python -m coverage report -m
+          python -m coverage json
+          export TOTAL=$(python -c "import json;print(json.load(open('coverage.json'))['totals']['percent_covered_display'])")
+          echo "total=$TOTAL" >> $GITHUB_ENV
+          if (( $TOTAL < 90 )); then
+            echo "Coverage is below 90%"
+            exit 1
+          fi
+        working-directory: ./lite
+      - name: run semantic segmentation tests and report coverage
+        run: |
+          pip install -e ".[test]"
+          COVERAGE_FILE=.coverage.segmentation python -m coverage run --omit "tests/*"  -m pytest -v tests/segmentation/
           python -m coverage combine
           python -m coverage report -m
           python -m coverage json

lite/tests/segmentation/conftest.py

@@ -0,0 +1,170 @@
+import numpy as np
+import pytest
+from valor_lite.segmentation import Bitmask, Segmentation
+
+
+def _generate_boolean_mask(
+    mask_shape: tuple[int, int],
+    annotation_shape: tuple[int, int, int, int],
+    label: str,
+) -> Bitmask:
+    mask = np.zeros(mask_shape, dtype=np.bool_)
+    xmin = annotation_shape[0]
+    xmax = annotation_shape[1]
+    ymin = annotation_shape[2]
+    ymax = annotation_shape[3]
+    mask[ymin:ymax, xmin:xmax] = True
+    return Bitmask(
+        mask=mask,
+        label=label,
+    )
+
+
+def _generate_random_boolean_mask(
+    mask_shape: tuple[int, int],
+    infill: float,
+    label: str,
+) -> Bitmask:
+    mask_size = mask_shape[0] * mask_shape[1]
+    mask = np.zeros(mask_size, dtype=np.bool_)
+
+    n_positives = int(mask_size * infill)
+    mask[:n_positives] = True
+    np.random.shuffle(mask)
+    mask = mask.reshape(mask_shape)
+
+    return Bitmask(
+        mask=mask,
+        label=label,
+    )
+
+
+@pytest.fixture
+def basic_segmentations() -> list[Segmentation]:
+
+    gmask1 = Bitmask(
+        mask=np.array([[True, False], [False, True]]),
+        label="v1",
+    )
+    gmask2 = Bitmask(
+        mask=np.array([[False, False], [True, False]]),
+        label="v2",
+    )
+
+    pmask1 = Bitmask(
+        mask=np.array([[True, False], [False, False]]),
+        label="v1",
+    )
+    pmask2 = Bitmask(
+        mask=np.array([[False, True], [True, False]]),
+        label="v2",
+    )
+
+    return [
+        Segmentation(
+            uid="uid0",
+            groundtruths=[gmask1, gmask2],
+            predictions=[pmask1, pmask2],
+        )
+    ]
+
+
+@pytest.fixture
+def segmentations_from_boxes() -> list[Segmentation]:
+
+    mask_shape = (900, 300)
+
+    rect1 = (0, 100, 0, 100)
+    rect2 = (150, 300, 400, 500)
+    rect3 = (50, 150, 0, 100)  # overlaps 50% with rect1
+    rect4 = (101, 151, 301, 401)  # overlaps 1 pixel with rect2
+
+    gmask1 = _generate_boolean_mask(mask_shape, rect1, "v1")
+    gmask2 = _generate_boolean_mask(mask_shape, rect2, "v2")
+
+    pmask1 = _generate_boolean_mask(mask_shape, rect3, "v1")
+    pmask2 = _generate_boolean_mask(mask_shape, rect4, "v2")
+
+    return [
+        Segmentation(
+            uid="uid1",
+            groundtruths=[gmask1],
+            predictions=[pmask1],
+        ),
+        Segmentation(
+            uid="uid2",
+            groundtruths=[gmask2],
+            predictions=[pmask2],
+        ),
+    ]
+
+
+@pytest.fixture
+def large_random_segmentations() -> list[Segmentation]:
+
+    mask_shape = (1000, 1000)
+    infills_per_seg = [
+        (0.9, 0.09, 0.01),
+        (0.4, 0.4, 0.1),
+        (0.3, 0.3, 0.3),
+    ]
+    labels_per_seg = [
+        ("v1", "v2", "v3"),
+        ("v4", "v5", "v6"),
+        ("v7", "v8", "v9"),
+    ]
+
+    return [
+        Segmentation(
+            uid=f"uid{idx}",
+            groundtruths=[
+                _generate_random_boolean_mask(mask_shape, infill, label)
+                for infill, label in zip(infills, labels)
+            ],
+            predictions=[
+                _generate_random_boolean_mask(mask_shape, infill, label)
+                for infill, label in zip(infills, labels)
+            ],
+        )
+        for idx, (infills, labels) in enumerate(
+            zip(infills_per_seg, labels_per_seg)
+        )
+    ]
+
+
+@pytest.fixture
+def massive_random_segmentations() -> list[Segmentation]:
+    """
+    A variant of `large_random_segmentations`.
+
+    This fixture is not used as it takes 10s to load.
+    """
+
+    mask_shape = (5000, 5000)
+    infills_per_seg = [
+        (0.9, 0.09, 0.01),
+        (0.4, 0.4, 0.1),
+        (0.3, 0.3, 0.3),
+    ]
+    labels_per_seg = [
+        ("v1", "v2", "v3"),
+        ("v4", "v5", "v6"),
+        ("v7", "v8", "v9"),
+    ]
+
+    return [
+        Segmentation(
+            uid=f"uid{idx}",
+            groundtruths=[
+                _generate_random_boolean_mask(mask_shape, infill, label)
+                for infill, label in zip(infills, labels)
+            ],
+            predictions=[
+                _generate_random_boolean_mask(mask_shape, infill, label)
+                for infill, label in zip(infills, labels)
+            ],
+        )
+        for idx, (infills, labels) in enumerate(
+            zip(infills_per_seg, labels_per_seg)
+        )
+    ]

lite/tests/segmentation/test_accuracy.py

@@ -0,0 +1,64 @@
+from valor_lite.segmentation import (
+    Accuracy,
+    DataLoader,
+    MetricType,
+    Segmentation,
+)
+
+
+def test_accuracy_basic_segmentations(basic_segmentations: list[Segmentation]):
+    loader = DataLoader()
+    loader.add_data(basic_segmentations)
+    evaluator = loader.finalize()
+
+    metrics = evaluator.evaluate(as_dict=True)
+
+    actual_metrics = [m for m in metrics[MetricType.Accuracy]]
+    expected_metrics = [
+        {
+            "type": "Accuracy",
+            "value": 0.5,
+            "parameters": {},
+        },
+    ]
+    for m in actual_metrics:
+        assert m in expected_metrics
+    for m in expected_metrics:
+        assert m in actual_metrics
+
+
+def test_accuracy_segmentations_from_boxes(
+    segmentations_from_boxes: list[Segmentation],
+):
+    loader = DataLoader()
+    loader.add_data(segmentations_from_boxes)
+    evaluator = loader.finalize()
+
+    metrics = evaluator.evaluate(as_dict=True)
+
+    actual_metrics = [m for m in metrics[MetricType.Accuracy]]
+    expected_metrics = [
+        {
+            "type": "Accuracy",
+            "value": 0.9444481481481481,
+            "parameters": {},
+        },
+    ]
+    for m in actual_metrics:
+        assert m in expected_metrics
+    for m in expected_metrics:
+        assert m in actual_metrics
+
+
+def test_accuracy_large_random_segmentations(
+    large_random_segmentations: list[Segmentation],
+):
+    loader = DataLoader()
+    loader.add_data(large_random_segmentations)
+    evaluator = loader.finalize()
+
+    metrics = evaluator.evaluate()[MetricType.Accuracy]
+
+    assert len(metrics) == 1
+    assert isinstance(metrics[0], Accuracy)
+    assert round(metrics[0].value, 1) == 0.5  # random choice

lite/tests/segmentation/test_annotation.py

@@ -0,0 +1,80 @@
+import numpy as np
+import pytest
+from valor_lite.segmentation import Bitmask, Segmentation
+
+
+def test_bitmask():
+
+    Bitmask(
+        mask=np.array([True, False]),
+        label="label",
+    )
+
+    with pytest.raises(ValueError) as e:
+        Bitmask(mask=np.array([1, 2]), label="label")
+    assert "int64" in str(e)
+
+
+def test_segmentation():
+
+    s = Segmentation(
+        uid="uid",
+        groundtruths=[
+            Bitmask(
+                mask=np.array([True, False]),
+                label="label",
+            )
+        ],
+        predictions=[
+            Bitmask(
+                mask=np.array([True, False]),
+                label="label",
+            )
+        ],
+    )
+    assert s.shape == (2,)
+    assert s.size == 2
+
+    with pytest.raises(ValueError) as e:
+        Segmentation(
+            uid="uid",
+            groundtruths=[
+                Bitmask(
+                    mask=np.array([True, False, False]),
+                    label="label",
+                )
+            ],
+            predictions=[
+                Bitmask(
+                    mask=np.array([True, False]),
+                    label="label",
+                )
+            ],
+        )
+    assert "mismatch" in str(e)
+
+    with pytest.raises(ValueError) as e:
+        Segmentation(
+            uid="uid",
+            groundtruths=[],
+            predictions=[
+                Bitmask(
+                    mask=np.array([True, False]),
+                    label="label",
+                )
+            ],
+        )
+    assert "missing ground truths" in str(e)
+
+    with pytest.raises(ValueError) as e:
+        Segmentation(
+            uid="uid",
+            groundtruths=[
+                Bitmask(
+                    mask=np.array([True, False]),
+                    label="label",
+                )
+            ],
+            predictions=[],
+        )
+    assert "missing predictions" in str(e)