diff --git a/cli/nu.py b/cli/nu.py index 70e1a4d2..a153465e 100644 --- a/cli/nu.py +++ b/cli/nu.py @@ -37,6 +37,15 @@ def nu(ctx, web): nu.add_command(slices) # type: ignore nu.add_command(tests) # type: ignore +# Local EvaluationV2 (`nu evalv2 run`). Guarded so `nu` still works if the optional evalv2-core +# dependency isn't installed yet (ships from a private index; see pyproject wiring). +try: + from nucleus.evalv2.cli import evalv2 + + nu.add_command(evalv2) # type: ignore +except ImportError: + pass + if __name__ == "__main__": """To debug, run this script followed by request command tree e.g. `cli/nu.py datasets list`""" nu() diff --git a/nucleus/evalv2/__init__.py b/nucleus/evalv2/__init__.py new file mode 100644 index 00000000..603d930f --- /dev/null +++ b/nucleus/evalv2/__init__.py @@ -0,0 +1,6 @@ +"""Local EvaluationV2 harness: run the shared evalv2-core kernel over parquet on a laptop. + +This subpackage is intentionally import-light — it does NOT import the heavy top-level ``nucleus`` +SDK. The parquet adapter depends only on ``evalv2_core`` + ``pyarrow``, mapping parquet columns +straight to the kernel's storage-agnostic types. +""" diff --git a/nucleus/evalv2/_tests/test_cli.py b/nucleus/evalv2/_tests/test_cli.py new file mode 100644 index 00000000..3637ff53 --- /dev/null +++ b/nucleus/evalv2/_tests/test_cli.py @@ -0,0 +1,138 @@ +"""Step 10 — run_local_eval + `nu evalv2 run` CLI, isolated from the heavy nucleus SDK. + +The evalv2 modules use proper absolute imports (`from nucleus.evalv2.X import ...`). To load them +without executing the heavy ``nucleus/__init__``, we register empty stub packages for ``nucleus`` and +``nucleus.evalv2`` in ``sys.modules``, then load each submodule from its file in dependency order. +""" + +import importlib.util +import json +import pathlib +import sys +import types + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from click.testing import CliRunner + +_EVALV2 = pathlib.Path(__file__).resolve().parent.parent + + +def _stub(name: str) -> None: + if name not in sys.modules: + m = types.ModuleType(name) + m.__path__ = [] # mark as package + sys.modules[name] = m + + +def _load(mod_name: str, filename: str): + _stub("nucleus") + _stub("nucleus.evalv2") + spec = importlib.util.spec_from_file_location(mod_name, _EVALV2 / filename) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod # register before exec so downstream imports resolve + spec.loader.exec_module(mod) + return mod + + +_load("nucleus.evalv2.parquet_io", "parquet_io.py") +run_mod = _load("nucleus.evalv2.run", "run.py") +cli_mod = _load("nucleus.evalv2.cli", "cli.py") + + +def _write_predictions(path, rows): + cols = {k: [r[k] for r in rows] for k in ["prediction_id", "dataset_item_id", "label", "confidence", "x", "y", "width", "height"]} + cols["metadata"] = [None] * len(rows) + pq.write_table(pa.table(cols), path) + + +def _write_gt(path, rows): + cols = {k: [r[k] for r in rows] for k in ["ground_truth_id", "dataset_item_id", "label", "x", "y", "width", "height"]} + cols["metadata"] = [None] * len(rows) + pq.write_table(pa.table(cols), path) + + +def _fixture(tmp_path): + preds = tmp_path / "p.parquet" + gt = tmp_path / "g.parquet" + _write_predictions( + preds, + [ + {"prediction_id": "p1", "dataset_item_id": "d1", "label": "car", "confidence": 0.9, "x": 0.0, "y": 0.0, "width": 20.0, "height": 20.0}, + {"prediction_id": "p2", "dataset_item_id": "d2", "label": "truck", "confidence": 0.3, "x": 0.0, "y": 0.0, "width": 20.0, "height": 20.0}, + ], + ) + _write_gt( + gt, + [ + {"ground_truth_id": "g1", "dataset_item_id": "d1", "label": "car", "x": 0.0, "y": 0.0, "width": 20.0, "height": 20.0}, + {"ground_truth_id": "g2", "dataset_item_id": "d2", "label": "truck", "x": 0.0, "y": 0.0, "width": 20.0, "height": 20.0}, + ], + ) + return preds, gt + + +# ---- config_from_dict ----------------------------------------------------------------------------- + + +def test_config_from_dict_parses_all() -> None: + cfg = run_mod.config_from_dict( + { + "rollup_groups": [{"class_name": "vehicle", "labels": ["car", "truck"]}], + "exclusion_rules": [{"type": "confidence", "min_confidence": 0.4}], + "class_agnostic": True, + "iou_type": "bbox", + "min_prediction_score": 0.2, + } + ) + assert cfg.rollup_groups[0].class_name == "vehicle" + assert cfg.exclusion_rules and cfg.exclusion_rules[0].type == "confidence" + assert cfg.class_agnostic is True and cfg.min_prediction_score == 0.2 + + +# ---- run_local_eval ------------------------------------------------------------------------------- + + +def test_run_local_eval_writes_bundle(tmp_path) -> None: + preds, gt = _fixture(tmp_path) + out = tmp_path / "out" + summary = run_mod.run_local_eval(preds, gt, out) + assert summary.total_gt == 2 + assert (out / "matches.parquet").exists() and (out / "summary.json").exists() + + +# ---- CLI ------------------------------------------------------------------------------------------ + + +def test_cli_run_smoke(tmp_path) -> None: + preds, gt = _fixture(tmp_path) + out = tmp_path / "out" + res = CliRunner().invoke(cli_mod.evalv2, ["run", "--predictions", str(preds), "--ground-truth", str(gt), "--out", str(out)]) + assert res.exit_code == 0, res.output + echoed = json.loads(res.output) + assert echoed["total_gt"] == 2 and echoed["out"] == str(out) + assert (out / "charts.json").exists() + + +def test_cli_config_and_flags_honored(tmp_path) -> None: + preds, gt = _fixture(tmp_path) + out = tmp_path / "out" + cfg = tmp_path / "cfg.json" + # rollup car+truck -> vehicle; without it d1(car) and d2(truck) each match same-label GTs anyway, + # so exercise min_prediction_score: drop p2 (0.3) via the flag -> d2 GT becomes an FN. + cfg.write_text(json.dumps({"rollup_groups": [{"class_name": "vehicle", "labels": ["car", "truck"]}]})) + res = CliRunner().invoke( + cli_mod.evalv2, + ["run", "--predictions", str(preds), "--ground-truth", str(gt), "--out", str(out), + "--config", str(cfg), "--min-prediction-score", "0.5"], + ) + assert res.exit_code == 0, res.output + echoed = json.loads(res.output) + # p1(0.9) kept -> TP; p2(0.3) dropped -> d2 GT is FN + assert echoed["total_tp"] == 1 and echoed["total_fn"] == 1 + + +def test_cli_missing_file_errors() -> None: + res = CliRunner().invoke(cli_mod.evalv2, ["run", "--predictions", "/nope.parquet", "--ground-truth", "/nope2.parquet", "--out", "/tmp/x"]) + assert res.exit_code != 0 diff --git a/nucleus/evalv2/_tests/test_parquet_io.py b/nucleus/evalv2/_tests/test_parquet_io.py new file mode 100644 index 00000000..1f11907b --- /dev/null +++ b/nucleus/evalv2/_tests/test_parquet_io.py @@ -0,0 +1,228 @@ +"""Round-trip test for the parquet adapter. + +Isolated from the heavy top-level ``nucleus`` package: parquet_io is loaded directly from its file +(so importing it never triggers ``nucleus/__init__``), and pytest is run with ``--confcutdir`` so the +repo-root ``conftest.py`` (which imports nucleus + requires an API key) is never collected. +""" + +import importlib.util +import json +import pathlib +from collections.abc import Iterable + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +import evalv2_core +from evalv2_core.types import ( + BoxGeometry, + EvalConfig, + EvalSummary, + GroundTruth, + ItemBundle, + Prediction, +) + +# --- load parquet_io.py directly, bypassing nucleus/__init__ -------------------------------------- +_PARQUET_IO = pathlib.Path(__file__).resolve().parent.parent / "parquet_io.py" +_spec = importlib.util.spec_from_file_location("evalv2_parquet_io", _PARQUET_IO) +_mod = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_mod) +ParquetSource = _mod.ParquetSource +ParquetSink = _mod.ParquetSink + + +# --- in-memory reference source/sink (for round-trip comparison) ----------------------------------- +class _MemSource: + def __init__(self, cfg, items): + self._cfg, self._items = cfg, items + + def read_config(self): + return self._cfg + + def read_items(self): + return iter(self._items) + + +class _MemSink: + def __init__(self): + self.summary = None + + def begin(self): + ... + + def write_matches(self, m, p): + ... + + def write_summary(self, summary, metrics, fields, charts): + self.summary = summary + + def finalize(self): + ... + + +def _write_predictions(path): + pq.write_table( + pa.table( + { + "prediction_id": ["p1", "p2"], + "dataset_item_id": ["d1", "d2"], + "label": ["car", "car"], + "confidence": [0.9, 0.8], + "x": [0.0, 0.0], + "y": [0.0, 0.0], + "width": [20.0, 20.0], + "height": [20.0, 20.0], + "metadata": [None, None], + } + ), + path, + ) + + +def _write_ground_truth(path): + pq.write_table( + pa.table( + { + "ground_truth_id": ["g1", "g2", "g3"], + "dataset_item_id": ["d1", "d2", "d3"], + "label": ["car", "person", "car"], + "x": [0.0, 0.0, 0.0], + "y": [0.0, 0.0, 0.0], + "width": [20.0, 20.0, 20.0], + "height": [20.0, 20.0, 20.0], + "metadata": [None, None, None], + } + ), + path, + ) + + +def _write_items(path): + pq.write_table( + pa.table( + { + "dataset_item_id": ["d1", "d2", "d3"], + "metadata": [json.dumps({"weather": "rain"}), json.dumps({"weather": "sun"}), json.dumps({"weather": "sun"})], + } + ), + path, + ) + + +def _reference_items(): + def box(*b): + return BoxGeometry(*b) + + return [ + ItemBundle("d1", [Prediction("p1", "d1", "car", 0.9, box(0, 0, 20, 20))], [GroundTruth("g1", "d1", "car", box(0, 0, 20, 20))], {"weather": "rain"}), + ItemBundle("d2", [Prediction("p2", "d2", "car", 0.8, box(0, 0, 20, 20))], [GroundTruth("g2", "d2", "person", box(0, 0, 20, 20))], {"weather": "sun"}), + ItemBundle("d3", [], [GroundTruth("g3", "d3", "car", box(0, 0, 20, 20))], {"weather": "sun"}), + ] + + +def test_parquet_round_trip_matches_in_memory(tmp_path) -> None: + pred_p = tmp_path / "predictions.parquet" + gt_p = tmp_path / "ground_truth.parquet" + items_p = tmp_path / "items.parquet" + out = tmp_path / "out" + _write_predictions(pred_p) + _write_ground_truth(gt_p) + _write_items(items_p) + + # run through the parquet adapter + src = ParquetSource(pred_p, gt_p, items_p, config=EvalConfig()) + sink = ParquetSink(out) + parquet_summary = evalv2_core.run(src, sink) + + # run the same items purely in memory + mem_sink = _MemSink() + evalv2_core.run(_MemSource(EvalConfig(), _reference_items()), mem_sink) + mem_summary = mem_sink.summary + + assert isinstance(parquet_summary, EvalSummary) + for key in ("map_50", "map_50_95", "total_tp", "total_fp", "total_fn", "total_gt"): + a, b = getattr(parquet_summary, key), getattr(mem_summary, key) + assert (a is None and b is None) or a == pytest.approx(b), key + + +def test_source_yields_joined_bundles(tmp_path) -> None: + _write_predictions(tmp_path / "p.parquet") + _write_ground_truth(tmp_path / "g.parquet") + _write_items(tmp_path / "i.parquet") + src = ParquetSource(tmp_path / "p.parquet", tmp_path / "g.parquet", tmp_path / "i.parquet") + items = list(src.read_items()) + assert [it.dataset_item_id for it in items] == ["d1", "d2", "d3"] + d1 = items[0] + assert d1.predictions[0].id == "p1" and d1.ground_truths[0].id == "g1" + assert d1.item_metadata == {"weather": "rain"} + assert isinstance(d1.predictions[0].geometry, BoxGeometry) + + +def test_output_bundle_files_written_and_readable(tmp_path) -> None: + _write_predictions(tmp_path / "p.parquet") + _write_ground_truth(tmp_path / "g.parquet") + out = tmp_path / "out" + src = ParquetSource(tmp_path / "p.parquet", tmp_path / "g.parquet") + evalv2_core.run(src, ParquetSink(out)) + + assert (out / "matches.parquet").exists() + assert (out / "per_threshold.parquet").exists() + matches = pq.read_table(out / "matches.parquet").to_pylist() + assert matches and {m["match_type"] for m in matches} <= {"TP", "FP", "FN"} + per_thr = pq.read_table(out / "per_threshold.parquet") + assert "size_bucket" in per_thr.column_names + + summary_doc = json.loads((out / "summary.json").read_text()) + assert "summary" in summary_doc and len(summary_doc["metrics_at_confidence"]) == 17 + charts = json.loads((out / "charts.json").read_text()) + for key in ("mapSummary", "prCurve", "confusionMatrix", "tideAttribution", "f1Curve"): + assert key in charts + + +def test_streaming_groups_across_batch_boundaries(tmp_path) -> None: + # force many tiny batches so an item's rows and item boundaries span batches + _write_predictions(tmp_path / "p.parquet") + _write_ground_truth(tmp_path / "g.parquet") + _write_items(tmp_path / "i.parquet") + src = ParquetSource( + tmp_path / "p.parquet", tmp_path / "g.parquet", tmp_path / "i.parquet", batch_size=1 + ) + items = list(src.read_items()) + assert [it.dataset_item_id for it in items] == ["d1", "d2", "d3"] + assert items[0].predictions[0].id == "p1" and items[0].ground_truths[0].id == "g1" + assert items[2].predictions == [] and items[2].ground_truths[0].id == "g3" # d3 gt-only + + +def test_unsorted_input_raises(tmp_path) -> None: + # predictions out of dataset_item_id order -> loud error, not a silent mis-join + pq.write_table( + pa.table( + { + "prediction_id": ["p2", "p1"], + "dataset_item_id": ["d2", "d1"], # descending -> unsorted + "label": ["car", "car"], + "confidence": [0.8, 0.9], + "x": [0.0, 0.0], "y": [0.0, 0.0], "width": [20.0, 20.0], "height": [20.0, 20.0], + "metadata": [None, None], + } + ), + tmp_path / "p.parquet", + ) + _write_ground_truth(tmp_path / "g.parquet") + src = ParquetSource(tmp_path / "p.parquet", tmp_path / "g.parquet", batch_size=1) + with pytest.raises(ValueError, match="sorted by dataset_item_id"): + list(src.read_items()) + + +def test_item_metadata_round_trips_into_match_rows(tmp_path) -> None: + _write_predictions(tmp_path / "p.parquet") + _write_ground_truth(tmp_path / "g.parquet") + _write_items(tmp_path / "i.parquet") + out = tmp_path / "out" + src = ParquetSource(tmp_path / "p.parquet", tmp_path / "g.parquet", tmp_path / "i.parquet") + evalv2_core.run(src, ParquetSink(out)) + matches = pq.read_table(out / "matches.parquet").to_pylist() + d1_rows = [m for m in matches if m["dataset_item_id"] == "d1"] + assert d1_rows and json.loads(d1_rows[0]["item_metadata"]) == {"weather": "rain"} diff --git a/nucleus/evalv2/cli.py b/nucleus/evalv2/cli.py new file mode 100644 index 00000000..a8e1e494 --- /dev/null +++ b/nucleus/evalv2/cli.py @@ -0,0 +1,73 @@ +"""`nu evalv2` CLI: run a local EvaluationV2 over parquet. + +Registered into the main ``nu`` CLI (cli/nu.py). Import-light; the heavy ``nucleus`` SDK is not +imported here. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import click + +from nucleus.evalv2.run import config_from_dict, run_local_eval + + +@click.group("evalv2") +def evalv2() -> None: + """Local EvaluationV2 over exported parquet (offline; upload separately).""" + + +@evalv2.command("run") +@click.option("--predictions", required=True, type=click.Path(exists=True), help="predictions.parquet") +@click.option("--ground-truth", required=True, type=click.Path(exists=True), help="ground_truth.parquet") +@click.option("--items", type=click.Path(exists=True), default=None, help="optional items.parquet (metadata)") +@click.option("--out", required=True, type=click.Path(), help="output directory for the result bundle") +@click.option("--config", "config_path", type=click.Path(exists=True), default=None, help="JSON EvalConfig") +@click.option("--iou-type", type=click.Choice(["bbox", "segm"]), default=None) +@click.option("--class-agnostic", is_flag=True, default=False) +@click.option("--min-prediction-score", type=float, default=None) +@click.option("--shard-size", type=int, default=2000, show_default=True) +def run_cmd( + predictions: str, + ground_truth: str, + items: str | None, + out: str, + config_path: str | None, + iou_type: str | None, + class_agnostic: bool, + min_prediction_score: float | None, + shard_size: int, +) -> None: + """Compute an evaluation locally and write matches/summary/charts to OUT.""" + cfg_dict = json.loads(Path(config_path).read_text()) if config_path else {} + if iou_type: + cfg_dict["iou_type"] = iou_type + if class_agnostic: + cfg_dict["class_agnostic"] = True + if min_prediction_score is not None: + cfg_dict["min_prediction_score"] = min_prediction_score + + summary = run_local_eval( + predictions, + ground_truth, + out, + items_path=items, + config=config_from_dict(cfg_dict), + shard_size=shard_size, + ) + click.echo( + json.dumps( + { + "map_50": summary.map_50, + "map_50_95": summary.map_50_95, + "total_tp": summary.total_tp, + "total_fp": summary.total_fp, + "total_fn": summary.total_fn, + "total_gt": summary.total_gt, + "out": str(out), + }, + indent=2, + ) + ) diff --git a/nucleus/evalv2/parquet_io.py b/nucleus/evalv2/parquet_io.py new file mode 100644 index 00000000..300629c4 --- /dev/null +++ b/nucleus/evalv2/parquet_io.py @@ -0,0 +1,302 @@ +"""Parquet adapter for the local (SDK) EvaluationV2 harness. + +Implements the evalv2-core ``EvalSource``/``EvalSink`` ports over parquet: + + - ``ParquetSource`` reads a predictions parquet + a ground-truth parquet (+ optional items parquet), + joins them on ``dataset_item_id``, and yields ``ItemBundle``s. Config is a constructor argument + (never stored in the parquet). + - ``ParquetSink`` streams match rows / per-threshold rows into parquet files and writes the + summary + charts as JSON — the result bundle the "upload eval" flow later pushes to the platform. + +Depends only on ``evalv2_core`` + ``pyarrow`` — never the heavy top-level ``nucleus`` SDK. + +Scale: ``ParquetSource`` streams via a sorted merge-join — it never loads a whole parquet into memory. +The input parquets MUST be sorted by ``dataset_item_id`` (string order); the source iterates all files +in row-group batches in lock-step, buffering only the current item's rows (memory is O(one item + +batch), independent of dataset size). Out-of-order input raises loudly rather than silently +mis-joining. + +Input schema (columns; each file sorted by dataset_item_id): + predictions.parquet : prediction_id, dataset_item_id, label, confidence, x, y, width, height, + [segmentation (JSON), metadata (JSON)] + ground_truth.parquet: ground_truth_id, dataset_item_id, label, x, y, width, height, + [segmentation (JSON), metadata (JSON)] + items.parquet (opt) : dataset_item_id, metadata (JSON) + +Output bundle (out_dir): + matches.parquet, per_threshold.parquet, summary.json, charts.json +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable, Iterator, Sequence +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import pyarrow as pa +import pyarrow.parquet as pq + +from evalv2_core.types import ( + BoxGeometry, + ChartsPayload, + EvalConfig, + EvalSummary, + Geometry, + GroundTruth, + ItemBundle, + MatchRow, + MetadataField, + MetricAtConfidence, + PerThresholdRow, + Prediction, + SegmentationGeometry, +) + +# ---- parquet output schemas ---------------------------------------------------------------------- + +_MATCH_SCHEMA = pa.schema( + [ + ("dataset_item_id", pa.string()), + ("match_type", pa.string()), + ("iou", pa.float64()), + ("true_positive", pa.bool_()), + ("model_prediction_id", pa.string()), + ("ground_truth_annotation_id", pa.string()), + ("pred_canonical_label", pa.string()), + ("gt_canonical_label", pa.string()), + ("pred_raw_label", pa.string()), + ("gt_raw_label", pa.string()), + ("confidence", pa.float64()), + ("gt_area", pa.float64()), + ("pred_area", pa.float64()), + ("item_metadata", pa.string()), # JSON + ("prediction_metadata", pa.string()), # JSON + ("nearest_same_class_gt_id", pa.string()), + ("nearest_same_class_gt_iou", pa.float64()), + ("nearest_any_class_gt_id", pa.string()), + ("nearest_any_class_gt_iou", pa.float64()), + ("nearest_any_class_gt_canonical_label", pa.string()), + ] +) + +_PER_THRESHOLD_SCHEMA = pa.schema( + [ + ("iou_threshold", pa.float64()), + ("dataset_item_id", pa.string()), + ("model_prediction_id", pa.string()), + ("pred_canonical_label", pa.string()), + ("confidence", pa.float64()), + ("matched_gt_id", pa.string()), + ("matched_gt_canonical_label", pa.string()), + ("iou", pa.float64()), + ("size_bucket", pa.string()), + ] +) + + +# ---- helpers ------------------------------------------------------------------------------------- + + +def _json_load(value: Any) -> dict: + if value is None or value == "": + return {} + if isinstance(value, dict): + return value + return json.loads(value) + + +def _geometry_from_row(row: dict) -> Geometry: + seg = row.get("segmentation") + if seg not in (None, ""): + return SegmentationGeometry(segmentation=seg if isinstance(seg, (list, dict)) else json.loads(seg)) + return BoxGeometry( + x=float(row["x"]), y=float(row["y"]), width=float(row["width"]), height=float(row["height"]) + ) + + +def _iter_item_groups(path: str, source_name: str, batch_size: int) -> Iterator[tuple[str, list[dict]]]: + """Yield (dataset_item_id, rows) for each contiguous block, streaming row-group batches. + + Requires the file sorted by ``dataset_item_id`` (string order). Buffers only the current item's + rows; raises if a group boundary goes backwards (unsorted input) to avoid silent mis-joins. + """ + current_id: str | None = None + buf: list[dict] = [] + for batch in pq.ParquetFile(path).iter_batches(batch_size=batch_size): + for row in batch.to_pylist(): + rid = row["dataset_item_id"] + if current_id is None: + current_id, buf = rid, [row] + elif rid == current_id: + buf.append(row) + else: + if rid < current_id: + raise ValueError( + f"{source_name} must be sorted by dataset_item_id " + f"(saw {rid!r} after {current_id!r})" + ) + yield current_id, buf + current_id, buf = rid, [row] + if current_id is not None: + yield current_id, buf + + +class _PeekableGroups: + """One-item lookahead over an ``_iter_item_groups`` generator.""" + + def __init__(self, gen: Iterator[tuple[str, list[dict]]]) -> None: + self._gen = gen + self._current: tuple[str, list[dict]] | None = None + self._advance() + + def _advance(self) -> None: + self._current = next(self._gen, None) + + @property + def current_id(self) -> str | None: + return self._current[0] if self._current is not None else None + + def take(self) -> list[dict]: + assert self._current is not None + rows = self._current[1] + self._advance() + return rows + + +# ---- source -------------------------------------------------------------------------------------- + + +class ParquetSource: + """EvalSource over two (or three) parquet files. Config is supplied here, not read from parquet.""" + + def __init__( + self, + predictions_path: str | Path, + ground_truth_path: str | Path, + items_path: str | Path | None = None, + config: EvalConfig | None = None, + *, + batch_size: int = 65_536, + ) -> None: + self._predictions_path = str(predictions_path) + self._ground_truth_path = str(ground_truth_path) + self._items_path = str(items_path) if items_path is not None else None + self._config = config or EvalConfig() + self._batch_size = batch_size + + def read_config(self) -> EvalConfig: + return self._config + + def read_items(self) -> Iterator[ItemBundle]: + """Streaming sorted merge-join over the (sorted) parquets. Memory is O(one item + batch).""" + preds = _PeekableGroups( + _iter_item_groups(self._predictions_path, "predictions.parquet", self._batch_size) + ) + gts = _PeekableGroups( + _iter_item_groups(self._ground_truth_path, "ground_truth.parquet", self._batch_size) + ) + items = _PeekableGroups( + _iter_item_groups(self._items_path, "items.parquet", self._batch_size) + if self._items_path is not None + else iter(()) + ) + + while True: + ids = [g.current_id for g in (preds, gts, items) if g.current_id is not None] + if not ids: + return + di = min(ids) # lexicographic — matches the sorted-string contract + pred_rows = preds.take() if preds.current_id == di else [] + gt_rows = gts.take() if gts.current_id == di else [] + meta_rows = items.take() if items.current_id == di else [] + yield ItemBundle( + dataset_item_id=di, + predictions=[ + Prediction( + id=r["prediction_id"], + dataset_item_id=di, + label=r["label"], + confidence=float(r["confidence"]), + geometry=_geometry_from_row(r), + metadata=_json_load(r.get("metadata")), + ) + for r in pred_rows + ], + ground_truths=[ + GroundTruth( + id=r["ground_truth_id"], + dataset_item_id=di, + label=r["label"], + geometry=_geometry_from_row(r), + metadata=_json_load(r.get("metadata")), + ) + for r in gt_rows + ], + item_metadata=_json_load(meta_rows[0].get("metadata")) if meta_rows else {}, + ) + + +# ---- sink ---------------------------------------------------------------------------------------- + + +def _match_row_to_dict(r: MatchRow) -> dict: + d = asdict(r) + d["item_metadata"] = json.dumps(d["item_metadata"]) + d["prediction_metadata"] = json.dumps(d["prediction_metadata"]) + return d + + +class ParquetSink: + """EvalSink writing the result bundle (matches/per_threshold parquet + summary/charts JSON).""" + + def __init__(self, out_dir: str | Path) -> None: + self.out_dir = Path(out_dir) + self._match_writer: pq.ParquetWriter | None = None + self._per_threshold_writer: pq.ParquetWriter | None = None + + def begin(self) -> None: + self.out_dir.mkdir(parents=True, exist_ok=True) + self._match_writer = pq.ParquetWriter(self.out_dir / "matches.parquet", _MATCH_SCHEMA) + self._per_threshold_writer = pq.ParquetWriter( + self.out_dir / "per_threshold.parquet", _PER_THRESHOLD_SCHEMA + ) + + def write_matches( + self, + match_rows: Iterable[MatchRow], + per_threshold_rows: Iterable[PerThresholdRow], + ) -> None: + assert self._match_writer is not None and self._per_threshold_writer is not None, "call begin()" + match_dicts = [_match_row_to_dict(r) for r in match_rows] + if match_dicts: + self._match_writer.write_table(pa.Table.from_pylist(match_dicts, schema=_MATCH_SCHEMA)) + per_thr_dicts = [asdict(r) for r in per_threshold_rows] + if per_thr_dicts: + self._per_threshold_writer.write_table( + pa.Table.from_pylist(per_thr_dicts, schema=_PER_THRESHOLD_SCHEMA) + ) + + def write_summary( + self, + summary: EvalSummary, + metrics_at_confidence: Sequence[MetricAtConfidence], + metadata_fields: Sequence[MetadataField], + charts: ChartsPayload, + ) -> None: + summary_doc = { + "summary": asdict(summary), + "metrics_at_confidence": [asdict(m) for m in metrics_at_confidence], + "metadata_fields": [asdict(f) for f in metadata_fields], + } + (self.out_dir / "summary.json").write_text(json.dumps(summary_doc, indent=2)) + (self.out_dir / "charts.json").write_text(json.dumps(dict(charts.payload), indent=2)) + + def finalize(self) -> None: + if self._match_writer is not None: + self._match_writer.close() + self._match_writer = None + if self._per_threshold_writer is not None: + self._per_threshold_writer.close() + self._per_threshold_writer = None diff --git a/nucleus/evalv2/run.py b/nucleus/evalv2/run.py new file mode 100644 index 00000000..f318f936 --- /dev/null +++ b/nucleus/evalv2/run.py @@ -0,0 +1,55 @@ +"""Programmatic entry for a local EvaluationV2 run: parquet in -> result bundle out. + +Wires the parquet adapter to the shared ``evalv2_core`` kernel. Import-light (``evalv2_core`` + +``parquet_io``); does not import the heavy top-level ``nucleus`` SDK. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from evalv2_core import run as core_run +from evalv2_core.exclusions import parse_exclusion_rules +from evalv2_core.types import EvalConfig, EvalSummary, RollupGroup + +from nucleus.evalv2.parquet_io import ParquetSink, ParquetSource + + +def config_from_dict(raw: dict[str, Any] | None) -> EvalConfig: + """Build an EvalConfig from a plain dict (CLI --config JSON). All keys optional.""" + d = raw or {} + kwargs: dict[str, Any] = {} + if d.get("rollup_groups"): + kwargs["rollup_groups"] = [ + RollupGroup(class_name=g["class_name"], labels=list(g["labels"])) + for g in d["rollup_groups"] + ] + if d.get("exclusion_rules"): + kwargs["exclusion_rules"] = parse_exclusion_rules(d["exclusion_rules"]) + if "class_agnostic" in d: + kwargs["class_agnostic"] = bool(d["class_agnostic"]) + if d.get("iou_type"): + kwargs["iou_type"] = d["iou_type"] + if d.get("min_prediction_score") is not None: + kwargs["min_prediction_score"] = float(d["min_prediction_score"]) + if d.get("confidence_grid"): + kwargs["confidence_grid"] = tuple(float(x) for x in d["confidence_grid"]) + return EvalConfig(**kwargs) + + +def run_local_eval( + predictions_path: str | Path, + ground_truth_path: str | Path, + out_dir: str | Path, + *, + items_path: str | Path | None = None, + config: EvalConfig | None = None, + shard_size: int = 2000, +) -> EvalSummary: + """Run an evaluation over local (sorted) parquets and write the result bundle to ``out_dir``.""" + source = ParquetSource( + predictions_path, ground_truth_path, items_path=items_path, config=config or EvalConfig() + ) + sink = ParquetSink(out_dir) + return core_run(source, sink, shard_size=shard_size)