wip cache upgrade

Author: czaloom

src/valor_lite/semantic_segmentation/evaluator.py

@@ -3,12 +3,15 @@
 from pathlib import Path
 
 import numpy as np
+import pyarrow as pa
 import pyarrow.compute as pc
 import pyarrow.dataset as ds
 from numpy.typing import NDArray
 
-from valor_lite.cache import DataType
+from valor_lite.cache import CacheReader, DataType
+from valor_lite.exceptions import EmptyCacheError
 from valor_lite.semantic_segmentation.computation import compute_metrics
+from valor_lite.semantic_segmentation.format import PathFormatter
 from valor_lite.semantic_segmentation.metric import Metric, MetricType
 from valor_lite.semantic_segmentation.utilities import (
     unpack_precision_recall_iou_into_metric_lists,
@@ -35,44 +38,171 @@ class Filter:
     predictions: pc.Expression | None = None
 
 
-class Evaluator:
+class Evaluator(PathFormatter):
     def __init__(
         self,
-        name: str = "default",
-        directory: str | Path = ".valor",
-        labels_override: dict[int, str] | None = None,
+        path: str | Path,
+        cache: CacheReader,
+        info: EvaluatorInfo,
+        index_to_label: dict[int, str],
+        confusion_matrix: NDArray[np.uint64],
     ):
-        self._directory = Path(directory)
-        self._name = name
-        self._path = self._directory / name
-        self._cache_path = self._path / "counts"
-        self._metadata_path = self._path / "metadata.json"
+        self._path = Path(path)
+        self._cache = cache
+        self._info = info
+        self._index_to_label = index_to_label
+        self._confusion_matrix = confusion_matrix
+
+    @classmethod
+    def load(
+        cls,
+        path: str | Path,
+        index_to_label_override: dict[int, str] | None = None,
+    ):
+        # validate path
+        path = Path(path)
+        if not path.exists():
+            raise FileNotFoundError(f"Directory does not exist: {path}")
+        elif not path.is_dir():
+            raise NotADirectoryError(
+                f"Path exists but is not a directory: {path}"
+            )
 
-        # link cache
-        self._dataset = ds.dataset(self._cache_path, format="parquet")
+        # load cache
+        cache = CacheReader.load(cls._generate_cache_path(path))
 
         # build evaluator meta
         (
-            self._index_to_label,
-            self._confusion_matrix,
-            self._info,
-        ) = self.generate_meta(self._dataset, labels_override)
+            index_to_label,
+            confusion_matrix,
+            info,
+        ) = cls.generate_meta(cache.dataset, index_to_label_override)
 
         # read config
-        with open(self._metadata_path, "r") as f:
+        metadata_path = cls._generate_metadata_path(path)
+        with open(metadata_path, "r") as f:
             types = json.load(f)
-            self._info.datum_metadata_types = types["datum"]
-            self._info.groundtruth_metadata_types = types["groundtruth"]
-            self._info.prediction_metadata_types = types["prediction"]
-        with open(self._cache_path / ".cfg", "r") as f:
-            cfg = json.load(f)
-            self._detailed_batch_size = cfg["batch_size"]
-            self._detailed_rows_per_file = cfg["rows_per_file"]
-            self._detailed_compression = cfg["compression"]
+            info.datum_metadata_types = types["datum_metadata_types"]
+            info.groundtruth_metadata_types = types[
+                "groundtruth_metadata_types"
+            ]
+            info.prediction_metadata_types = types["prediction_metadata_types"]
+
+        return cls(
+            path=path,
+            cache=cache,
+            info=info,
+            index_to_label=index_to_label,
+            confusion_matrix=confusion_matrix,
+        )
+
+    def filter(
+        self,
+        path: str | Path,
+        filter_expr: Filter,
+    ) -> "Evaluator":
+        """
+        Filter evaluator cache.
+
+        Parameters
+        ----------
+        path : str | Path
+            Where to store the filtered cache.
+        filter_expr : Filter
+            An object containing filter expressions.
+
+        Returns
+        -------
+        Evaluator
+            A new evaluator object containing the filtered cache.
+        """
+        from valor_lite.semantic_segmentation.loader import Loader
+
+        loader = Loader.create(
+            path=path,
+            batch_size=self.cache.batch_size,
+            rows_per_file=self.cache.rows_per_file,
+            compression=self.cache.compression,
+            datum_metadata_types=self.info.datum_metadata_types,
+            groundtruth_metadata_types=self.info.groundtruth_metadata_types,
+            prediction_metadata_types=self.info.prediction_metadata_types,
+        )
+        for fragment in self.cache.dataset.get_fragments():
+            tbl = fragment.to_table(filter=filter_expr.datums)
+
+            columns = (
+                "datum_id",
+                "gt_label_id",
+                "pd_label_id",
+            )
+            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)
+
+            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_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_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[~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])
+                )
+            loader._cache.write_table(tbl)
+
+        loader._cache.flush()
+        if loader._cache.dataset.count_rows() == 0:
+            raise EmptyCacheError()
+
+        return loader.finalize()
+
+    def delete(self):
+        """
+        Delete evaluator cache.
+        """
+        from valor_lite.semantic_segmentation.loader import Loader
+
+        Loader.delete(self.path)
+
+    @property
+    def path(self) -> Path:
+        return self._path
 
     @property
-    def dataset(self) -> ds.Dataset:
-        return self._dataset
+    def cache(self) -> CacheReader:
+        return self._cache
 
     @property
     def info(self) -> EvaluatorInfo:
@@ -185,105 +315,6 @@ def generate_meta(
 
         return labels, matrix, info
 
-    def filter(
-        self,
-        filter_expr: Filter,
-        name: str | None = None,
-        directory: str | Path | None = None,
-    ) -> "Evaluator":
-        """
-        Filter evaluator cache.
-
-        Parameters
-        ----------
-        filter_expr : Filter
-            An object containing filter expressions.
-        name : str, optional
-            Filtered cache name.
-        directory : str | Path, optional
-            The directory to store the filtered cache.
-
-        Returns
-        -------
-        Evaluator
-            A new evaluator object containing the filtered cache.
-        """
-        loader = cls(
-            directory=directory,
-            name=name,
-            batch_size=evaluator._detailed_batch_size,
-            rows_per_file=evaluator._detailed_rows_per_file,
-            compression=evaluator._detailed_compression,
-            datum_metadata_types=evaluator.info.datum_metadata_types,
-            groundtruth_metadata_types=evaluator.info.groundtruth_metadata_types,
-            prediction_metadata_types=evaluator.info.prediction_metadata_types,
-        )
-        for fragment in evaluator.dataset.get_fragments():
-            tbl = fragment.to_table(filter=filter_expr.datums)
-
-            columns = (
-                "datum_id",
-                "gt_label_id",
-                "pd_label_id",
-            )
-            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)
-
-            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_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_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[~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])
-                )
-            loader._cache.write_table(tbl)
-
-        loader._cache.flush()
-        if loader._cache.dataset.count_rows() == 0:
-            raise EmptyCacheError()
-
-        evaluator = Evaluator(
-            directory=loader._directory,
-            name=loader._name,
-            labels_override=evaluator._index_to_label,
-        )
-        return evaluator
-
     def compute_precision_recall_iou(self) -> dict[MetricType, list]:
         """
         Performs an evaluation and returns metrics.