evaluation_utils.py

from __future__ import annotations

import csv
import json
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, Iterable, List


KNOWN_SUPPORT_IMAGENET_MAP: Dict[str, str] = {
    "airplane": "404_airliner",
    "backpack": "414_backpack, back pack, knapsack, packsack, rucksack, haversack",
    "banana": "954_banana",
    "bear": "295_American black bear, black bear, Ursus americanus, Euarctos americanus",
    "bicycle": "671_mountain bike, all-terrain bike, off-roader",
    "bird": "094_hummingbird",
    "bottle": "898_water bottle",
    "broccoli": "937_broccoli",
    "cat": "281_tabby, tabby cat",
    "chair": "559_folding chair",
    "clock": "409_analog clock",
    "couch": "831_studio couch, day bed",
    "cup": "968_cup",
    "dog": "245_French bulldog",
    "elephant": "385_Indian elephant, Elephas maximus",
    "keyboard": "508_computer keyboard, keypad",
    "laptop": "620_laptop, laptop computer",
    "orange": "950_orange",
    "pizza": "963_pizza, pizza pie",
    "tv": "851_television, television system",
}

UNKNOWN_EVAL_IMAGENET_MAP: Dict[str, str] = {
    "bench": "703_park bench",
    "boat": "814_speedboat",
    "bus": "779_school bus",
    "cell_phone": "487_cellular telephone, cellular phone, cellphone, cell, mobile phone",
    "teddy_bear": "850_teddy, teddy bear",
    "truck": "717_pickup, pickup truck",
    "vase": "883_vase",
    "zebra": "340_zebra",
}


def write_json(path: str | Path, payload: Any) -> None:
    out_path = Path(path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")


def write_csv(path: str | Path, rows: Iterable[Dict[str, Any]]) -> None:
    out_path = Path(path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    rows = list(rows)
    if not rows:
        out_path.write_text("", encoding="utf-8")
        return

    fieldnames: List[str] = []
    seen = set()
    for row in rows:
        for key in row.keys():
            if key not in seen:
                seen.add(key)
                fieldnames.append(key)

    with out_path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            serialized = {
                key: json.dumps(value, sort_keys=True) if isinstance(value, (dict, list)) else value
                for key, value in row.items()
            }
            writer.writerow(serialized)


def build_stream_manifest(items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    known = [dict(item) for item in items if item.get("eval_split") == "known"]
    unknown = [dict(item) for item in items if item.get("eval_split") == "unknown"]

    def _group(entries: List[Dict[str, Any]], label_key: str) -> Dict[int, List[Dict[str, Any]]]:
        grouped: Dict[int, List[Dict[str, Any]]] = defaultdict(list)
        for entry in entries:
            grouped[int(entry.get("round_index", 0))].append(entry)
        for round_index in grouped:
            grouped[round_index].sort(key=lambda item: str(item.get(label_key, "")))
        return grouped

    known_rounds = _group(known, "expected_label")
    unknown_rounds = _group(unknown, "source_label")
    round_indices = sorted(set(known_rounds.keys()) | set(unknown_rounds.keys()))

    ordered: List[Dict[str, Any]] = []
    for round_index in round_indices:
        ordered.extend(known_rounds.get(round_index, []))
        ordered.extend(unknown_rounds.get(round_index, []))

    for stream_index, entry in enumerate(ordered):
        entry["stream_index"] = stream_index
    return ordered


def summarize_evaluation_records(records: List[Dict[str, Any]]) -> Dict[str, Any]:
    total = len(records)
    known = [row for row in records if row.get("eval_split") == "known"]
    unknown = [row for row in records if row.get("eval_split") == "unknown"]

    def _rate(numerator: int, denominator: int) -> float:
        return float(numerator) / float(denominator) if denominator else 0.0

    known_top1 = sum(1 for row in known if bool(row.get("known_correct_top1", False)))
    known_any = sum(1 for row in known if bool(row.get("known_correct_any", False)))
    unknown_rejected = sum(1 for row in unknown if bool(row.get("unknown_rejected", False)))
    unknown_false_accept = sum(1 for row in unknown if bool(row.get("unknown_false_accept", False)))
    no_detections = sum(1 for row in records if int(row.get("num_detections", 0)) == 0)

    accepted_updates_total = sum(int(row.get("accepted_updates_count", 0)) for row in records)
    accepted_updates_known_correct = sum(
        int(row.get("accepted_updates_count", 0))
        for row in known
        if bool(row.get("known_correct_top1", False))
    )
    accepted_updates_known_incorrect = sum(
        int(row.get("accepted_updates_count", 0))
        for row in known
        if not bool(row.get("known_correct_top1", False))
    )
    accepted_updates_unknown = sum(int(row.get("accepted_updates_count", 0)) for row in unknown)

    known_by_round: Dict[int, List[Dict[str, Any]]] = defaultdict(list)
    unknown_by_round: Dict[int, List[Dict[str, Any]]] = defaultdict(list)
    for row in known:
        known_by_round[int(row.get("round_index", 0))].append(row)
    for row in unknown:
        unknown_by_round[int(row.get("round_index", 0))].append(row)

    known_round_accuracy = {
        str(round_index): _rate(
            sum(1 for row in rows if bool(row.get("known_correct_any", False))),
            len(rows),
        )
        for round_index, rows in sorted(known_by_round.items())
    }
    unknown_round_rejection = {
        str(round_index): _rate(
            sum(1 for row in rows if bool(row.get("unknown_rejected", False))),
            len(rows),
        )
        for round_index, rows in sorted(unknown_by_round.items())
    }

    final_support_total = 0
    initial_support_total = 0
    if records:
        initial_support_total = int(records[0].get("memory_support_total_before", 0))
        final_support_total = int(records[-1].get("memory_support_total_after", initial_support_total))

    growth_by_class: Dict[str, int] = {}
    if records:
        first = records[0].get("memory_class_counts_before", {}) or {}
        last = records[-1].get("memory_class_counts_after", {}) or {}
        all_labels = sorted(set(first.keys()) | set(last.keys()))
        growth_by_class = {
            label: int(last.get(label, 0)) - int(first.get(label, 0))
            for label in all_labels
        }

    return {
        "total_images": total,
        "known_images": len(known),
        "unknown_images": len(unknown),
        "no_detection_rate": _rate(no_detections, total),
        "known_top1_accuracy": _rate(known_top1, len(known)),
        "known_any_match_accuracy": _rate(known_any, len(known)),
        "unknown_rejection_rate": _rate(unknown_rejected, len(unknown)),
        "unknown_false_accept_rate": _rate(unknown_false_accept, len(unknown)),
        "overall_image_score": _rate(known_any + unknown_rejected, total),
        "accepted_updates_total": accepted_updates_total,
        "accepted_updates_known_correct": accepted_updates_known_correct,
        "accepted_updates_known_incorrect": accepted_updates_known_incorrect,
        "accepted_updates_unknown": accepted_updates_unknown,
        "initial_support_total": initial_support_total,
        "final_support_total": final_support_total,
        "support_growth_total": final_support_total - initial_support_total,
        "support_growth_by_class": growth_by_class,
        "known_any_match_accuracy_by_round": known_round_accuracy,
        "unknown_rejection_rate_by_round": unknown_round_rejection,
    }


def build_comparison(static_summary: Dict[str, Any], dynamic_summary: Dict[str, Any]) -> Dict[str, Any]:
    keys = [
        "known_top1_accuracy",
        "known_any_match_accuracy",
        "unknown_rejection_rate",
        "unknown_false_accept_rate",
        "overall_image_score",
        "accepted_updates_total",
        "support_growth_total",
    ]
    delta = {}
    for key in keys:
        static_value = static_summary.get(key, 0.0)
        dynamic_value = dynamic_summary.get(key, 0.0)
        if isinstance(static_value, (int, float)) and isinstance(dynamic_value, (int, float)):
            delta[key] = float(dynamic_value) - float(static_value)
    return {
        "static": static_summary,
        "dynamic": dynamic_summary,
        "delta_dynamic_minus_static": delta,
    }