evaluate_dfod.py

from __future__ import annotations

import argparse
import copy
import json
import shutil
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List

from Baseline_Perception.embedder import clear_embedder_cache
from Baseline_Reasoning.serialization import load_memory_bundle, save_memory
from build_memory_from_support import build_memory_bundle
from dfod_runtime import build_runtime_config, run_dfod_inference
from evaluation_utils import build_comparison, summarize_evaluation_records, write_csv, write_json


def _load_manifest(path: str) -> List[Dict[str, Any]]:
    return json.loads(Path(path).read_text(encoding="utf-8"))


def _clone_adapter_weights(source_path: str, target_path: str) -> str:
    src = Path(source_path)
    dst = Path(target_path)
    dst.parent.mkdir(parents=True, exist_ok=True)
    if src.exists():
        shutil.copy2(src, dst)
    return str(dst)


def _top_detection(detections: List[Dict[str, Any]]) -> Dict[str, Any] | None:
    if not detections:
        return None
    return detections[0]


def _build_eval_record(
    sample: Dict[str, Any],
    run_result: Dict[str, Any],
    mode: str,
) -> Dict[str, Any]:
    detections = list(run_result.get("detections", []))
    accepted_updates = list(run_result.get("accepted_updates", []))
    top = _top_detection(detections)

    expected_label = str(sample.get("expected_label", "unknown"))
    known_correct_any = any(
        (not bool(det.get("pred_unknown", True))) and str(det.get("pred_label", "")) == expected_label
        for det in detections
    )
    top_label = str(top.get("pred_label")) if top is not None and top.get("pred_label") is not None else None
    top_unknown = bool(top.get("pred_unknown", True)) if top is not None else True
    known_correct_top1 = bool(
        sample.get("eval_split") == "known"
        and top is not None
        and not top_unknown
        and top_label == expected_label
    )
    unknown_rejected = bool(
        sample.get("eval_split") == "unknown"
        and (not detections or all(bool(det.get("pred_unknown", True)) for det in detections))
    )
    unknown_false_accept = bool(
        sample.get("eval_split") == "unknown"
        and any(not bool(det.get("pred_unknown", True)) for det in detections)
    )

    return {
        "mode": mode,
        "stream_index": int(sample.get("stream_index", -1)),
        "round_index": int(sample.get("round_index", 0)),
        "eval_split": str(sample.get("eval_split", "")),
        "expected_label": expected_label,
        "source_label": str(sample.get("source_label", "")),
        "image_path": str(sample.get("image_path", "")),
        "num_detections": len(detections),
        "top_pred_label": top_label,
        "top_pred_unknown": top_unknown,
        "top_pred_score": float(top.get("pred_score", 0.0)) if top is not None else None,
        "top_pred_margin": float(top.get("pred_margin", 0.0)) if top is not None else None,
        "top_rejection_reason": str(top.get("rejection_reason", "no_detection")) if top is not None else "no_detection",
        "known_correct_top1": known_correct_top1,
        "known_correct_any": known_correct_any,
        "unknown_rejected": unknown_rejected,
        "unknown_false_accept": unknown_false_accept,
        "accepted_updates_count": len(accepted_updates),
        "accepted_update_labels": [str(item.get("label", "")) for item in accepted_updates],
        "memory_support_total_before": int(run_result.get("memory_support_total_before", 0)),
        "memory_support_total_after": int(run_result.get("memory_support_total_after", 0)),
        "memory_class_counts_before": run_result.get("memory_class_counts_before", {}),
        "memory_class_counts_after": run_result.get("memory_class_counts_after", {}),
    }


def _run_mode(
    manifest: List[Dict[str, Any]],
    memory_bundle: Dict[str, Any],
    config: Dict[str, Any],
    mode: str,
    device: str,
) -> Dict[str, Any]:
    current_memory = copy.deepcopy(memory_bundle)
    records: List[Dict[str, Any]] = []

    for sample in manifest:
        run_result = run_dfod_inference(
            image_path=str(sample["image_path"]),
            memory_bundle=current_memory,
            config=config,
            device=device,
            apply_updates_enabled=bool(config["updates"].get("enabled", True)),
        )
        current_memory = run_result["memory_bundle"]
        records.append(_build_eval_record(sample=sample, run_result=run_result, mode=mode))

    summary = summarize_evaluation_records(records)
    return {
        "records": records,
        "summary": summary,
        "final_memory_bundle": current_memory,
    }


def evaluate_dfod(
    manifest_path: str,
    memory_path: str = "memory.npz",
    output_dir: str | None = None,
    device: str = "cpu",
    rebuild_memory: bool = False,
    max_images: int | None = None,
    calibration_path: str | None = None,
) -> Dict[str, Any]:
    manifest = _load_manifest(manifest_path)
    if max_images is not None:
        manifest = manifest[:max(0, int(max_images))]
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    out_dir = Path(output_dir) if output_dir else Path("evaluation_runs") / timestamp
    out_dir.mkdir(parents=True, exist_ok=True)

    if rebuild_memory:
        rebuilt_bundle = build_memory_bundle()
        save_memory(rebuilt_bundle, memory_path)

    base_memory = load_memory_bundle(memory_path)
    base_config = build_runtime_config({
        "updates": {"persist_updated_memory": False},
    })

    base_adapter_path = str(base_config["adapter"].get("weights_path", "adapter.pt"))
    static_adapter_path = _clone_adapter_weights(base_adapter_path, out_dir / "adapter_static.pt")
    dynamic_adapter_path = _clone_adapter_weights(base_adapter_path, out_dir / "adapter_dynamic.pt")

    static_config = build_runtime_config({
        "calibration": {
            "enabled": bool(calibration_path),
            "calibration_path": calibration_path,
        },
        "adapter": {"weights_path": static_adapter_path},
        "updates": {"enabled": False, "persist_updated_memory": False},
    })
    dynamic_config = build_runtime_config({
        "calibration": {
            "enabled": bool(calibration_path),
            "calibration_path": calibration_path,
        },
        "adapter": {"weights_path": dynamic_adapter_path},
        "updates": {"enabled": True, "persist_updated_memory": False},
    })

    clear_embedder_cache()
    static_result = _run_mode(
        manifest=manifest,
        memory_bundle=base_memory,
        config=static_config,
        mode="static",
        device=device,
    )

    clear_embedder_cache()
    dynamic_result = _run_mode(
        manifest=manifest,
        memory_bundle=base_memory,
        config=dynamic_config,
        mode="dynamic",
        device=device,
    )
    clear_embedder_cache()

    write_json(out_dir / "static_summary.json", static_result["summary"])
    write_json(out_dir / "dynamic_summary.json", dynamic_result["summary"])
    write_json(out_dir / "comparison.json", build_comparison(static_result["summary"], dynamic_result["summary"]))
    write_json(out_dir / "static_records.json", static_result["records"])
    write_json(out_dir / "dynamic_records.json", dynamic_result["records"])
    write_csv(out_dir / "static_records.csv", static_result["records"])
    write_csv(out_dir / "dynamic_records.csv", dynamic_result["records"])
    save_memory(dynamic_result["final_memory_bundle"], str(out_dir / "dynamic_final_memory.npz"))

    return {
        "output_dir": str(out_dir),
        "static_summary": static_result["summary"],
        "dynamic_summary": dynamic_result["summary"],
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Run a static-vs-dynamic DFOD evaluation stream.")
    parser.add_argument("--manifest", default="test_images/manifest.json")
    parser.add_argument("--memory-path", default="memory.npz")
    parser.add_argument("--output-dir", default=None)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--rebuild-memory", action="store_true")
    parser.add_argument("--max-images", type=int, default=None)
    parser.add_argument("--calibration-path", default=None)
    args = parser.parse_args()

    result = evaluate_dfod(
        manifest_path=args.manifest,
        memory_path=args.memory_path,
        output_dir=args.output_dir,
        device=args.device,
        rebuild_memory=args.rebuild_memory,
        max_images=args.max_images,
        calibration_path=args.calibration_path,
    )

    print(f"Saved evaluation outputs to {result['output_dir']}")
    print(
        "Static known-any={:.3f}, dynamic known-any={:.3f}".format(
            float(result["static_summary"]["known_any_match_accuracy"]),
            float(result["dynamic_summary"]["known_any_match_accuracy"]),
        )
    )
    print(
        "Static unknown-reject={:.3f}, dynamic unknown-reject={:.3f}".format(
            float(result["static_summary"]["unknown_rejection_rate"]),
            float(result["dynamic_summary"]["unknown_rejection_rate"]),
        )
    )


if __name__ == "__main__":
    main()