added filtering tests

Author: czaloom

src/valor_lite/cache.py

@@ -60,7 +60,7 @@ def convert_type_mapping_to_schema(
     """
     if not type_mapping:
         return []
-    return [(k, v.to_arrow()) for k, v in type_mapping.items()]
+    return [(k, DataType(v).to_arrow()) for k, v in type_mapping.items()]
 
 
 class CacheReader:

src/valor_lite/semantic_segmentation/loader.py

@@ -350,6 +350,8 @@ def filter(
             pairs[~mask_valid_gt, 1] = -1
             pairs[~mask_valid_pd, 2] = -1
 
+            print(pairs)
+
             for idx, col in enumerate(columns):
                 tbl = tbl.set_column(
                     tbl.schema.names.index(col), col, pa.array(pairs[:, idx])

tests/semantic_segmentation/conftest.py

@@ -18,6 +18,10 @@ def _generate_boolean_mask(
     return Bitmask(
         mask=mask,
         label=label,
+        metadata={
+            "gt_xmin": xmin,
+            "pd_xmin": xmin,
+        },
     )
 
 

tests/semantic_segmentation/test_filtering.py

@@ -1,11 +1,26 @@
 import numpy as np
+import pyarrow.compute as pc
 import pytest
 
+from valor_lite.cache import DataType
 from valor_lite.exceptions import EmptyFilterError
 from valor_lite.semantic_segmentation import DataLoader, Segmentation
+from valor_lite.semantic_segmentation.evaluator import Filter
 
 
-def test_filtering(segmentations_from_boxes: list[Segmentation]):
+def test_filtering_raises(segmentations_from_boxes: list[Segmentation]):
+
+    loader = DataLoader()
+    loader.add_data(segmentations_from_boxes)
+    evaluator = loader.finalize()
+    assert evaluator._confusion_matrix.shape == (3, 3)
+
+    with pytest.raises(EmptyFilterError):
+        evaluator.create_filter(datums=[])
+    assert evaluator._confusion_matrix.shape == (3, 3)
+
+
+def test_filtering_by_datum(segmentations_from_boxes: list[Segmentation]):
 
     loader = DataLoader()
     loader.add_data(segmentations_from_boxes)
@@ -51,13 +66,105 @@ def test_filtering(segmentations_from_boxes: list[Segmentation]):
         evaluator.create_filter(datums=[])
 
 
-def test_filtering_raises(segmentations_from_boxes: list[Segmentation]):
+def test_filtering_by_annotation_metadata(
+    segmentations_from_boxes: list[Segmentation],
+):
 
-    loader = DataLoader()
+    loader = DataLoader(
+        groundtruth_metadata_types={
+            "gt_xmin": DataType.FLOAT,
+        },
+        prediction_metadata_types={
+            "pd_xmin": DataType.FLOAT,
+        },
+    )
     loader.add_data(segmentations_from_boxes)
     evaluator = loader.finalize()
-    assert evaluator._confusion_matrix.shape == (3, 3)
 
-    with pytest.raises(EmptyFilterError):
-        evaluator.create_filter(datums=[])
-    assert evaluator._confusion_matrix.shape == (3, 3)
+    total_pixels = 540_000
+    assert evaluator.metadata.number_of_datums == 2
+    assert evaluator.metadata.number_of_labels == 2
+    assert evaluator.metadata.number_of_ground_truths == 25000
+    assert evaluator.metadata.number_of_predictions == 15000
+    assert evaluator.metadata.number_of_pixels == total_pixels
+
+    # test groundtruth filtering
+    filter_ = Filter(groundtruths=pc.field("gt_xmin") < 100)
+    filtered_evaluator = evaluator.filter(filter_)
+    confusion_matrix = filtered_evaluator._confusion_matrix
+    assert np.all(
+        confusion_matrix
+        == np.array(
+            [
+                [520000, 5000, 5000],
+                [5000, 5000, 0],
+                [0, 0, 0],
+            ]
+        )
+    )
+    assert confusion_matrix.sum() == total_pixels
+
+    filter_ = Filter(groundtruths=pc.field("gt_xmin") > 100)
+    filtered_evaluator = evaluator.filter(filter_)
+    confusion_matrix = filtered_evaluator._confusion_matrix
+    assert np.all(
+        confusion_matrix
+        == np.array(
+            [
+                [510001, 10000, 4999],
+                [0, 0, 0],
+                [14999, 0, 1],
+            ]
+        )
+    )
+    assert confusion_matrix.sum() == total_pixels
+
+    # test prediction filtering
+    filter_ = Filter(predictions=pc.field("pd_xmin") < 100)
+    filtered_evaluator = evaluator.filter(filter_)
+    confusion_matrix = filtered_evaluator._confusion_matrix
+    assert np.all(
+        confusion_matrix
+        == np.array(
+            [
+                [510000, 5000, 0],
+                [5000, 5000, 0],
+                [15000, 0, 0],
+            ]
+        )
+    )
+    assert confusion_matrix.sum() == total_pixels
+
+    filter_ = Filter(predictions=pc.field("pd_xmin") > 100)
+    filtered_evaluator = evaluator.filter(filter_)
+    confusion_matrix = filtered_evaluator._confusion_matrix
+    assert np.all(
+        confusion_matrix
+        == np.array(
+            [
+                [510001, 0, 4999],
+                [10000, 0, 0],
+                [14999, 0, 1],
+            ]
+        )
+    )
+    assert confusion_matrix.sum() == total_pixels
+
+    # filter out all gts and pds
+    filter_ = Filter(
+        groundtruths=pc.field("gt_xmin") > 1000,
+        predictions=pc.field("pd_xmin") > 1000,
+    )
+    filtered_evaluator = evaluator.filter(filter_)
+    confusion_matrix = filtered_evaluator._confusion_matrix
+    assert np.all(
+        confusion_matrix
+        == np.array(
+            [
+                [total_pixels, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+            ]
+        )
+    )
+    assert confusion_matrix.sum() == total_pixels