calibrate_dfod.py

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

from Baseline_Reasoning.calibration import build_openworld_calibration_artifact, save_calibration
from dfod_runtime import build_runtime_config, run_dfod_inference
from evaluation_utils import write_csv, write_json


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


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


def _calibration_record(sample: Dict[str, Any], run_result: Dict[str, Any]) -> Dict[str, Any]:
    detections = list(run_result.get("detections", []))
    top = _top_detection(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
    return {
        "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": str(sample.get("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": bool(
            sample.get("eval_split") == "known"
            and top is not None
            and not top_unknown
            and top_label == str(sample.get("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))
        ),
    }


def calibrate_dfod(
    manifest_path: str,
    memory_path: str = "memory.npz",
    output_path: str = "calibration.npz",
    device: str = "cpu",
) -> Dict[str, Any]:
    manifest = _load_manifest(manifest_path)
    config = build_runtime_config({
        "calibration": {"enabled": False},
        "reasoning": {
            "score_threshold": float("-inf"),
            "margin_threshold": float("-inf"),
        },
        "updates": {"enabled": False, "persist_updated_memory": False},
    })

    records: List[Dict[str, Any]] = []
    for sample in manifest:
        run_result = run_dfod_inference(
            image_path=str(sample["image_path"]),
            memory_path=memory_path,
            config=config,
            device=device,
            apply_updates_enabled=False,
        )
        records.append(_calibration_record(sample=sample, run_result=run_result))

    artifact = build_openworld_calibration_artifact(records=records, config=config)
    artifact["manifest_path"] = str(Path(manifest_path).resolve())
    artifact["memory_path"] = str(Path(memory_path).resolve())

    out_path = Path(output_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    save_calibration(artifact, str(out_path))
    write_json(out_path.with_suffix(".json"), artifact)
    write_json(out_path.with_name(f"{out_path.stem}_records.json"), records)
    write_csv(out_path.with_name(f"{out_path.stem}_records.csv"), records)
    return {
        "artifact": artifact,
        "output_path": str(out_path),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Calibrate DFOD rejection and update thresholds from a held-out stream.")
    parser.add_argument("--manifest", default="calibration_images/manifest.json")
    parser.add_argument("--memory-path", default="memory.npz")
    parser.add_argument("--output-path", default="calibration.npz")
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()

    result = calibrate_dfod(
        manifest_path=args.manifest,
        memory_path=args.memory_path,
        output_path=args.output_path,
        device=args.device,
    )
    artifact = result["artifact"]
    print(f"Saved calibration artifact to {result['output_path']}")
    print(
        "Rejection score={:.6f} margin={:.6f} known={:.3f} unknown_reject={:.3f}".format(
            float(artifact["score_threshold"]),
            float(artifact["margin_threshold"]),
            float(artifact["known_rate"]),
            float(artifact["unknown_rejection_rate"]),
        )
    )
    print(
        "Update score={:.6f} margin={:.6f} precision={:.3f} recall={:.3f}".format(
            float(artifact["update_score_threshold"]),
            float(artifact["update_margin_threshold"]),
            float(artifact["update_precision"]),
            float(artifact["update_recall"]),
        )
    )


if __name__ == "__main__":
    main()