| from __future__ import annotations |
|
|
| import argparse |
| from datetime import datetime |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| from sklearn.metrics import confusion_matrix |
|
|
| from .ijepa_localization import IJepaPatchLocalizer |
| from .obstacle_dataset import DEFAULT_OBSTACLE_DATASET, load_balanced_obstacle_rows, load_obstacle_image, parse_yolo_boxes |
| from .prototypes import build_class_prototypes, guess_objects_with_prototypes |
| from .small_head import guess_objects_with_head, train_small_head |
|
|
|
|
| def run_bulk_eval( |
| dataset_name: str, |
| split: str, |
| model_name: str, |
| eval_samples: int, |
| support_samples: int, |
| seed: int, |
| output: str, |
| ) -> Path: |
| run_id = datetime.now().strftime("%Y%m%d_%H%M%S") |
| output_path = resolve_output_path(output, run_id) |
| localizer = IJepaPatchLocalizer(model_name=model_name) |
| support_rows = load_balanced_obstacle_rows( |
| dataset_name, |
| split, |
| support_samples, |
| random_seed=seed + 10_000, |
| ) |
| support_files = {row["file_name"] for row in support_rows} |
| eval_rows = load_disjoint_eval_rows(dataset_name, split, eval_samples, seed, support_files) |
| prototypes = build_class_prototypes(dataset_name, split, support_rows, localizer) |
| head = train_small_head(dataset_name, split, support_rows, localizer) |
|
|
| records = [] |
| for sample_index, row in enumerate(eval_rows): |
| image = load_obstacle_image(dataset_name, row, split) |
| boxes = parse_yolo_boxes(row) |
| prototype_guesses = guess_objects_with_prototypes(image, boxes, localizer, prototypes) |
| head_guesses = guess_objects_with_head(image, boxes, localizer, head) |
| head_by_object = {guess.object_index: guess for guess in head_guesses} |
| prototype_by_object = {guess.object_index: guess for guess in prototype_guesses} |
|
|
| for object_index, box in enumerate(boxes, start=1): |
| prototype_guess = prototype_by_object.get(object_index) |
| head_guess = head_by_object.get(object_index) |
| records.append( |
| { |
| "sample": sample_index, |
| "run_id": run_id, |
| "support_samples": support_samples, |
| "eval_samples": eval_samples, |
| "seed": seed, |
| "file_name": row["file_name"], |
| "object": object_index, |
| "yolo_label": box.class_name, |
| "prototype_guess": prototype_guess.ijepa_guess if prototype_guess else None, |
| "prototype_similarity": prototype_guess.similarity if prototype_guess else None, |
| "prototype_agreement": prototype_guess.agreement if prototype_guess else None, |
| "head_guess": head_guess.head_guess if head_guess else None, |
| "head_confidence": head_guess.confidence if head_guess else None, |
| "head_agreement": head_guess.agreement if head_guess else None, |
| "head_train_objects": head.train_objects if head else None, |
| "head_train_accuracy": head.train_accuracy if head else None, |
| "head_parameter_count": head.parameter_count if head else None, |
| "prototype_classes": len(prototypes), |
| } |
| ) |
|
|
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| df = pd.DataFrame(records) |
| df.to_csv(output_path, index=False) |
| write_summaries(df, output_path) |
| return output_path |
|
|
|
|
| def load_disjoint_eval_rows(dataset_name: str, split: str, eval_samples: int, seed: int, support_files: set[str]): |
| candidates = load_balanced_obstacle_rows( |
| dataset_name, |
| split, |
| eval_samples + len(support_files) + 100, |
| random_seed=seed, |
| ) |
| selected = [] |
| for row in candidates: |
| if row["file_name"] in support_files: |
| continue |
| selected.append(row) |
| if len(selected) >= eval_samples: |
| break |
| if len(selected) < eval_samples: |
| raise ValueError(f"Could only find {len(selected)} disjoint eval rows; requested {eval_samples}.") |
| return selected |
|
|
|
|
| def write_summaries(df: pd.DataFrame, output_path: Path) -> None: |
| if df.empty: |
| return |
|
|
| summary = { |
| "run_id": str(df["run_id"].iloc[0]) if "run_id" in df else None, |
| "support_samples": first_int(df, "support_samples"), |
| "eval_samples": first_int(df, "eval_samples"), |
| "seed": first_int(df, "seed"), |
| "objects": int(len(df)), |
| "classes": int(df["yolo_label"].nunique()), |
| "prototype_accuracy": safe_accuracy(df, "prototype_agreement"), |
| "head_accuracy": safe_accuracy(df, "head_agreement"), |
| "prototype_macro_accuracy": macro_accuracy(df, "prototype_agreement"), |
| "head_macro_accuracy": macro_accuracy(df, "head_agreement"), |
| "head_train_objects": first_int(df, "head_train_objects"), |
| "head_parameter_count": first_int(df, "head_parameter_count"), |
| "prototype_classes": first_int(df, "prototype_classes"), |
| } |
| summary_path = sibling_path(output_path, "_summary.csv") |
| per_class_path = sibling_path(output_path, "_per_class.csv") |
| prototype_confusion_path = sibling_path(output_path, "_prototype_confusion.csv") |
| head_confusion_path = sibling_path(output_path, "_head_confusion.csv") |
| report_path = sibling_path(output_path, "_report.json") |
|
|
| pd.DataFrame([summary]).to_csv(summary_path, index=False) |
|
|
| per_class = ( |
| df.groupby("yolo_label") |
| .agg( |
| objects=("yolo_label", "size"), |
| prototype_accuracy=("prototype_agreement", safe_mean), |
| head_accuracy=("head_agreement", safe_mean), |
| avg_prototype_similarity=("prototype_similarity", "mean"), |
| avg_head_confidence=("head_confidence", "mean"), |
| ) |
| .reset_index() |
| .sort_values(["head_accuracy", "objects"], ascending=[True, False]) |
| ) |
| per_class.to_csv(per_class_path, index=False) |
|
|
| prototype_confusion = confusion_frame(df["yolo_label"], df["prototype_guess"]) |
| head_confusion = confusion_frame(df["yolo_label"], df["head_guess"]) |
| prototype_confusion.to_csv(prototype_confusion_path) |
| head_confusion.to_csv(head_confusion_path) |
| report = { |
| "summary": summary, |
| "per_class": per_class.replace({np.nan: None}).to_dict(orient="records"), |
| "top_prototype_confusions": top_confusions(df, "prototype_guess"), |
| "top_head_confusions": top_confusions(df, "head_guess"), |
| "files": { |
| "objects": str(output_path), |
| "summary": str(summary_path), |
| "per_class": str(per_class_path), |
| "prototype_confusion": str(prototype_confusion_path), |
| "head_confusion": str(head_confusion_path), |
| }, |
| } |
| report_path.write_text(json.dumps(report, indent=2), encoding="utf-8") |
|
|
|
|
| def sibling_path(output_path: Path, suffix: str) -> Path: |
| names = { |
| "_summary.csv": "summary.csv", |
| "_per_class.csv": "per_class.csv", |
| "_prototype_confusion.csv": "prototype_confusion.csv", |
| "_head_confusion.csv": "head_confusion.csv", |
| "_report.json": "report.json", |
| } |
| return output_path.with_name(names.get(suffix, f"{output_path.stem}{suffix}")) |
|
|
|
|
| def resolve_output_path(output: str, run_id: str) -> Path: |
| path = Path(output) |
| if output == "auto": |
| return Path("outputs") / f"run_{run_id}" / "objects.csv" |
| if "{timestamp}" in output or "{run_id}" in output: |
| resolved = Path(output.format(timestamp=run_id, run_id=run_id)) |
| if resolved.suffix: |
| return resolved |
| return resolved / "objects.csv" |
| if path.suffix == "": |
| return path / "objects.csv" |
| if path.exists(): |
| return path.with_name(f"{path.stem}_{run_id}{path.suffix}") |
| return path |
|
|
|
|
| def safe_mean(values) -> float: |
| clean = pd.Series(values).dropna() |
| if clean.empty: |
| return np.nan |
| return float(clean.astype(bool).mean()) |
|
|
|
|
| def safe_accuracy(df: pd.DataFrame, column: str) -> float | None: |
| clean = df[column].dropna() |
| if clean.empty: |
| return None |
| return float(clean.astype(bool).mean()) |
|
|
|
|
| def macro_accuracy(df: pd.DataFrame, column: str) -> float | None: |
| clean = df.dropna(subset=[column]) |
| if clean.empty: |
| return None |
| return float(clean.groupby("yolo_label")[column].apply(safe_mean).mean()) |
|
|
|
|
| def confusion_frame(y_true, y_pred) -> pd.DataFrame: |
| clean = pd.DataFrame({"true": y_true, "pred": y_pred}).dropna() |
| labels = sorted(set(clean["true"]).union(set(clean["pred"]))) |
| matrix = confusion_matrix(clean["true"], clean["pred"], labels=labels) |
| return pd.DataFrame(matrix, index=labels, columns=labels) |
|
|
|
|
| def top_confusions(df: pd.DataFrame, prediction_column: str, limit: int = 10) -> list[dict]: |
| clean = df.dropna(subset=[prediction_column]) |
| wrong = clean[clean["yolo_label"] != clean[prediction_column]] |
| if wrong.empty: |
| return [] |
| counts = ( |
| wrong.groupby(["yolo_label", prediction_column]) |
| .size() |
| .reset_index(name="count") |
| .sort_values("count", ascending=False) |
| .head(limit) |
| ) |
| return [ |
| {"yolo_label": row["yolo_label"], "predicted": row[prediction_column], "count": int(row["count"])} |
| for _, row in counts.iterrows() |
| ] |
|
|
|
|
| def first_int(df: pd.DataFrame, column: str) -> int | None: |
| clean = df[column].dropna() |
| if clean.empty: |
| return None |
| return int(clean.iloc[0]) |
|
|
|
|
| def build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description="Run bulk I-JEPA prototype/head evaluation.") |
| parser.add_argument("--dataset-name", default=DEFAULT_OBSTACLE_DATASET) |
| parser.add_argument("--split", default="train") |
| parser.add_argument("--model-name", default="facebook/ijepa_vith14_1k") |
| parser.add_argument("--eval-samples", type=int, default=50) |
| parser.add_argument( |
| "--support-samples", |
| "--reference-samples", |
| dest="support_samples", |
| type=int, |
| default=80, |
| help="Images used to build class prototypes and train the tiny classifier.", |
| ) |
| parser.add_argument("--seed", type=int, default=7) |
| parser.add_argument( |
| "--output", |
| default="auto", |
| help="Output CSV path or run directory. Supports {timestamp} or {run_id}. Default: outputs/run_<timestamp>/objects.csv", |
| ) |
| return parser |
|
|
|
|
| def main() -> None: |
| args = build_parser().parse_args() |
| output = run_bulk_eval(**vars(args)) |
| print(f"Saved bulk evaluation: {output}") |
| for suffix in [ |
| "_summary.csv", |
| "_per_class.csv", |
| "_prototype_confusion.csv", |
| "_head_confusion.csv", |
| "_report.json", |
| ]: |
| print(f"Saved: {sibling_path(output, suffix)}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
