passing tests

Author: czaloom

src/valor_lite/semantic_segmentation/computation.py

@@ -109,7 +109,7 @@ def compute_intermediates(
 
 
 def compute_metrics(
-    confusion_matrices: NDArray[np.int64],
+    counts: NDArray[np.uint64],
 ) -> tuple[
     NDArray[np.float64],
     NDArray[np.float64],
@@ -126,8 +126,8 @@ def compute_metrics(
 
     Parameters
     ----------
-    confusion_matrices : NDArray[np.int64]
-        A 3-D array containing confusion matrices for each datum with shape (n_datums, n_labels + 1, n_labels + 1).
+    counts : NDArray[np.int64]
+        A 2-D confusion matrix with shape (n_labels + 1, n_labels + 1).
     label_metadata : NDArray[np.int64]
         A 2-D array containing label metadata with shape (n_labels, 2).
             Index 0: Ground Truth Label Count
@@ -150,15 +150,10 @@ def compute_metrics(
     NDArray[np.float64]
         Unmatched ground truth ratios.
     """
-    n_labels = confusion_matrices.shape[-1] - 1
-    n_pixels = confusion_matrices.sum()
-    label_metadata = np.zeros((n_labels, 2), dtype=np.int64)
-    label_metadata[:, 0] = confusion_matrices[:, 1:, :].sum(axis=(0, 2))
-    label_metadata[:, 1] = confusion_matrices[:, :, 1:].sum(axis=(0, 1))
-    gt_counts = label_metadata[:, 0]
-    pd_counts = label_metadata[:, 1]
-
-    counts = confusion_matrices.sum(axis=0)
+    n_labels = counts.shape[0] - 1
+    n_pixels = counts.sum()
+    gt_counts = counts[1:, :].sum(axis=1)
+    pd_counts = counts[:, 1:].sum(axis=0)
 
     # compute iou, unmatched_ground_truth and unmatched predictions
     intersection_ = counts[1:, 1:]

src/valor_lite/semantic_segmentation/evaluator.py

@@ -1,28 +1,14 @@
 import json
-from collections import defaultdict
-from dataclasses import asdict, dataclass
+from dataclasses import dataclass
 from pathlib import Path
 
 import numpy as np
+import pyarrow.compute as pc
+import pyarrow.dataset as ds
 from numpy.typing import NDArray
-from pyarrow import pa
-from pyarrow.compute import pc
-from pyarrow.dataset import ds
-from tqdm import tqdm
-
-from valor_lite.cache import (
-    CacheReader,
-    DataType,
-    convert_type_mapping_to_schema,
-)
-from valor_lite.exceptions import EmptyCacheError, EmptyFilterError
-from valor_lite.semantic_segmentation.annotation import Segmentation
-from valor_lite.semantic_segmentation.computation import (
-    compute_intermediates,
-    compute_label_metadata,
-    compute_metrics,
-    filter_cache,
-)
+
+from valor_lite.cache import DataType
+from valor_lite.semantic_segmentation.computation import compute_metrics
 from valor_lite.semantic_segmentation.metric import Metric, MetricType
 from valor_lite.semantic_segmentation.utilities import (
     unpack_precision_recall_iou_into_metric_lists,
@@ -31,12 +17,12 @@
 
 @dataclass
 class EvaluatorInfo:
+    number_of_rows: int = 0
     number_of_datums: int = 0
     number_of_labels: int = 0
     number_of_pixels: int = 0
     number_of_groundtruth_pixels: int = 0
     number_of_prediction_pixels: int = 0
-    number_of_rows: int = 0
     datum_metadata_types: dict[str, DataType] | None = None
     groundtruth_metadata_types: dict[str, DataType] | None = None
     prediction_metadata_types: dict[str, DataType] | None = None
@@ -68,7 +54,7 @@ def __init__(
         # build evaluator meta
         (
             self._index_to_label,
-            self._number_of_groundtruths_per_label,
+            self._confusion_matrix,
             self._info,
         ) = self.generate_meta(self._dataset, labels_override)
 
@@ -111,12 +97,11 @@ def generate_meta(
         -------
         labels : dict[int, str]
             Mapping of label ID's to label values.
-        number_of_groundtruths_per_label : NDArray[np.uint64]
-            Array of size (n_labels,) containing ground truth counts.
+        confusion_matrix : NDArray[np.uint64]
+            Array of size (n_labels + 1, n_labels + 1) containing pair counts.
         info : EvaluatorInfo
             Evaluator cache details.
         """
-        gt_counts_per_lbl = defaultdict(int)
         labels = labels_override if labels_override else {}
         info = EvaluatorInfo()
 
@@ -126,7 +111,7 @@ def generate_meta(
                 "datum_id",
                 "gt_label_id",
                 "pd_label_id",
-                "counts",
+                "count",
             )
             ids = np.column_stack(
                 [tbl[col].to_numpy() for col in columns]
@@ -159,21 +144,9 @@ def generate_meta(
             pd_labels.pop(-1, None)
             labels.update(pd_labels)
 
-            # count gts per label
-            gts = ids[:, 1].astype(np.int64)
-            unique_ann = np.unique(gts[gts[:, 0] >= 0], axis=0)
-            unique_labels, label_counts = np.unique(
-                unique_ann[:, 1], return_counts=True
-            )
-            for label_id, count in zip(unique_labels, label_counts):
-                gt_counts_per_lbl[int(label_id)] += int(count)
-
         # post-process
         labels.pop(-1, None)
 
-        # complete info object
-        info.number_of_labels = len(labels)
-
         # create confusion matrix
         n_labels = len(labels)
         matrix = np.zeros((n_labels + 1, n_labels + 1), dtype=np.uint64)
@@ -187,8 +160,11 @@ def generate_meta(
             ids = np.column_stack(
                 [tbl[col].to_numpy() for col in columns]
             ).astype(np.int64)
-            counts = tbl["counts"].to_numpy()
+            counts = tbl["count"].to_numpy()
 
+            mask_null_gts = ids[:, 1] == -1
+            mask_null_pds = ids[:, 2] == -1
+            matrix[0, 0] = counts[mask_null_gts & mask_null_pds].sum()
             for idx in range(n_labels):
                 mask_gts = ids[:, 1] == idx
                 for pidx in range(n_labels):
@@ -197,35 +173,18 @@ def generate_meta(
                         mask_gts & mask_pds
                     ].sum()
 
-                mask_unmatched_gts = mask_gts & (ids[:, 2] == -1)
+                mask_unmatched_gts = mask_gts & mask_null_pds
                 matrix[idx + 1, 0] = counts[mask_unmatched_gts].sum()
-                mask_unmatched_pds = (ids[:, 1] == -1) & (ids[:, 2] == idx)
-                matrix[0, idx + 1] = counts[mask_unmatched_pds]
+                mask_unmatched_pds = mask_null_gts & (ids[:, 2] == idx)
+                matrix[0, idx + 1] = counts[mask_unmatched_pds].sum()
 
-        return labels, matrix, info
-
-    @staticmethod
-    def iterate_pairs(
-        dataset: ds.Dataset,
-        columns: list[str] | None = None,
-    ):
-        for fragment in dataset.get_fragments():
-            tbl = fragment.to_table(columns=columns)
-            yield np.column_stack(
-                [tbl.column(i).to_numpy() for i in range(tbl.num_columns)]
-            )
+        # complete info object
+        info.number_of_labels = len(labels)
+        info.number_of_pixels = matrix.sum()
+        info.number_of_groundtruth_pixels = matrix[1:, :].sum()
+        info.number_of_prediction_pixels = matrix[:, 1:].sum()
 
-    @staticmethod
-    def iterate_pairs_with_table(
-        dataset: ds.Dataset,
-        columns: list[str] | None = None,
-    ):
-        for fragment in dataset.get_fragments():
-            tbl = fragment.to_table()
-            columns = columns if columns else tbl.columns
-            yield tbl, np.column_stack(
-                [tbl[col].to_numpy() for col in columns]
-            )
+        return labels, matrix, info
 
     def filter(
         self,
@@ -260,3 +219,29 @@ def filter(
             evaluator=self,
             filter_expr=filter_expr,
         )
+
+    def compute_precision_recall_iou(self) -> dict[MetricType, list]:
+        """
+        Performs an evaluation and returns metrics.
+
+        Returns
+        -------
+        dict[MetricType, list]
+            A dictionary mapping MetricType enumerations to lists of computed metrics.
+        """
+        results = compute_metrics(counts=self._confusion_matrix)
+        return unpack_precision_recall_iou_into_metric_lists(
+            results=results,
+            index_to_label=self._index_to_label,
+        )
+
+    def evaluate(self) -> dict[MetricType, list[Metric]]:
+        """
+        Computes all available metrics.
+
+        Returns
+        -------
+        dict[MetricType, list[Metric]]
+            Lists of metrics organized by metric type.
+        """
+        return self.compute_precision_recall_iou()

src/valor_lite/semantic_segmentation/loader.py

@@ -39,9 +39,6 @@ def __init__(
         self._labels: dict[str, int] = {}
         self._index_to_label: dict[int, str] = {}
         self._datum_count = 0
-        self._groundtruth_pixel_count = 0
-        self._prediction_pixel_count = 0
-        self._total_pixel_count = 0
 
         with open(self._metadata_path, "w") as f:
             types = {
@@ -247,18 +244,28 @@ def add_data(
                         },
                     ]
                 )
+            rows.append(
+                {
+                    # datum
+                    "datum_uid": segmentation.uid,
+                    "datum_id": self._datum_count,
+                    **datum_metadata,
+                    # groundtruth
+                    "gt_label": None,
+                    "gt_label_id": -1,
+                    # prediction
+                    "pd_label": None,
+                    "pd_label_id": -1,
+                    # pair
+                    "count": counts[0, 0],
+                }
+            )
             self._cache.write_rows(rows)
 
             # update datum cache
             self._datum_count += 1
 
-    def finalize(
-        self,
-        rows_per_file: int | None = None,
-        compression: str | None = None,
-        write_batch_size: int | None = None,
-        read_batch_size: int = 1000,
-    ):
+    def finalize(self):
         """
         Performs data finalization and some preprocessing steps.
 
@@ -271,18 +278,10 @@ def finalize(
         if self._cache.dataset.count_rows() == 0:
             raise EmptyCacheError()
 
-        evaluator = Evaluator(
+        return Evaluator(
             directory=self._directory,
             name=self._name,
         )
-        evaluator.rank(
-            where=self._ranked_path,
-            rows_per_file=rows_per_file,
-            compression=compression,
-            write_batch_size=write_batch_size,
-            read_batch_size=read_batch_size,
-        )
-        return evaluator
 
     @classmethod
     def filter(
@@ -302,59 +301,60 @@ def filter(
             groundtruth_metadata_types=evaluator.info.groundtruth_metadata_types,
             prediction_metadata_types=evaluator.info.prediction_metadata_types,
         )
-        for fragment in evaluator.detailed.get_fragments():
+        for fragment in evaluator.dataset.get_fragments():
             tbl = fragment.to_table(filter=filter_expr.datums)
 
             columns = (
                 "datum_id",
-                "gt_id",
-                "pd_id",
                 "gt_label_id",
                 "pd_label_id",
-                "iou",
-                "score",
             )
             pairs = np.column_stack([tbl[col].to_numpy() for col in columns])
 
             n_pairs = pairs.shape[0]
             gt_ids = pairs[:, (0, 1)].astype(np.int64)
             pd_ids = pairs[:, (0, 2)].astype(np.int64)
 
-            mask_valid_gt = np.zeros(n_pairs, dtype=np.bool_)
-            mask_valid_pd = np.zeros(n_pairs, dtype=np.bool_)
-
             if filter_expr.groundtruths is not None:
+                mask_valid_gt = np.zeros(n_pairs, dtype=np.bool_)
                 gt_tbl = tbl.filter(filter_expr.groundtruths)
                 gt_pairs = np.column_stack(
-                    [gt_tbl[col].to_numpy() for col in ("datum_id", "gt_id")]
+                    [
+                        gt_tbl[col].to_numpy()
+                        for col in ("datum_id", "gt_label_id")
+                    ]
                 ).astype(np.int64)
                 for gt in np.unique(gt_pairs, axis=0):
                     mask_valid_gt |= (gt_ids == gt).all(axis=1)
+            else:
+                mask_valid_gt = np.ones(n_pairs, dtype=np.bool_)
 
             if filter_expr.predictions is not None:
+                mask_valid_pd = np.zeros(n_pairs, dtype=np.bool_)
                 pd_tbl = tbl.filter(filter_expr.predictions)
                 pd_pairs = np.column_stack(
-                    [pd_tbl[col].to_numpy() for col in ("datum_id", "pd_id")]
+                    [
+                        pd_tbl[col].to_numpy()
+                        for col in ("datum_id", "pd_label_id")
+                    ]
                 ).astype(np.int64)
                 for pd in np.unique(pd_pairs, axis=0):
                     mask_valid_pd |= (pd_ids == pd).all(axis=1)
+            else:
+                mask_valid_pd = np.ones(n_pairs, dtype=np.bool_)
 
             mask_valid = mask_valid_gt | mask_valid_pd
             mask_valid_gt &= mask_valid
             mask_valid_pd &= mask_valid
 
-            pairs[np.ix_(~mask_valid_gt, (1, 3))] = -1.0  # type: ignore - numpy ix_
-            pairs[np.ix_(~mask_valid_pd, (2, 4, 6))] = -1.0  # type: ignore - numpy ix_
-            pairs[~mask_valid_pd | ~mask_valid_gt, 5] = 0.0
+            pairs[~mask_valid_gt, 1] = -1
+            pairs[~mask_valid_pd, 2] = -1
 
             for idx, col in enumerate(columns):
                 tbl = tbl.set_column(
                     tbl.schema.names.index(col), col, pa.array(pairs[:, idx])
                 )
-
-            mask_invalid = ~mask_valid | (pairs[:, (1, 2)] < 0).all(axis=1)
-            filtered_tbl = tbl.filter(pa.array(~mask_invalid))
-            loader._cache.write_table(filtered_tbl)
+            loader._cache.write_table(tbl)
 
         loader._cache.flush()
         if loader._cache.dataset.count_rows() == 0:
@@ -365,5 +365,4 @@ def filter(
             name=loader._name,
             labels_override=evaluator._index_to_label,
         )
-        evaluator.rank(where=loader._ranked_path)
         return evaluator