scripts/correctness.py

#!/usr/bin/env python3
"""Numerical equivalence: candidate ONNX vs FP32 baseline on identical inputs.

Feeds the SAME app-faithful `images` tensor to both graphs (via Python onnxruntime, CPU EP),
then reports both raw-output error and end-to-end annotation agreement. The end-to-end metrics
mirror the app's edgecrafter-seg parser + mask decode, so a small logit error near the
maskThreshold=0.0 boundary that flips a pixel is actually counted, not averaged away.

NOTE ON SCOPE: this measures *numerical* agreement against FP32 using Python ORT. It is the
decision metric for "can the candidate replace FP32" (FP32 is the reference). It is NOT the
browser latency benchmark — that is a separate harness (bench_browser). A candidate that agrees
here still must LOAD and RUN in ort-web WASM, which the browser harness verifies.

Usage:
  python correctness.py --baseline FP32.onnx --candidate CAND.onnx \
      --images DIR [DIR ...] [--limit N] [--json OUT.json]
"""

from __future__ import annotations

import argparse
import json
import os
import sys

import numpy as np
import onnxruntime as ort

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import ecseg_common as ec  # noqa: E402


def make_session(path: str) -> ort.InferenceSession:
    so = ort.SessionOptions()
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    so.intra_op_num_threads = 1
    so.inter_op_num_threads = 1
    return ort.InferenceSession(path, sess_options=so, providers=["CPUExecutionProvider"])


def run(session: ort.InferenceSession, x: np.ndarray) -> dict:
    names = [o.name for o in session.get_outputs()]
    outs = session.run(names, {"images": x})
    return dict(zip(names, outs))


def pct(x) -> float:
    return round(float(x) * 100, 3)


def compare_image(base_out: dict, cand_out: dict) -> dict:
    """Raw-output error + end-to-end agreement for one image."""
    bl = base_out["labels"].reshape(-1)
    cl = cand_out["labels"].reshape(-1)
    bs = base_out["scores"].reshape(-1).astype(np.float64)
    cs = cand_out["scores"].reshape(-1).astype(np.float64)
    bb = base_out["boxes"].reshape(-1, 4).astype(np.float64)
    cb = cand_out["boxes"].reshape(-1, 4).astype(np.float64)
    bm = base_out["masks"].astype(np.float32)  # [1,300,160,160]
    cm = cand_out["masks"].astype(np.float32)

    # ---- raw output errors (all 300 queries) ----
    label_agree = float(np.mean(bl == cl))
    score_mae = float(np.mean(np.abs(bs - cs)))
    score_max = float(np.max(np.abs(bs - cs)))
    box_mae = float(np.mean(np.abs(bb - cb)))
    box_max = float(np.max(np.abs(bb - cb)))
    mask_logit_mae = float(np.mean(np.abs(bm - cm)))
    mask_logit_max = float(np.max(np.abs(bm - cm)))

    # fraction of mask pixels whose binary decision at logit>0 flips
    b_bin = bm > ec.MASK_THRESHOLD
    c_bin = cm > ec.MASK_THRESHOLD
    mask_flip_frac = float(np.mean(b_bin != c_bin))

    # NaN / Inf hygiene on the candidate
    nan_inf = bool(
        np.isnan(cs).any() or np.isinf(cs).any() or
        np.isnan(cb).any() or np.isinf(cb).any() or
        np.isnan(cm).any() or np.isinf(cm).any()
    )

    # near-threshold sensitivity: queries whose baseline score is within ±0.05 of 0.4
    near = np.abs(bs - ec.CONF_THRESHOLD) <= 0.05
    near_count = int(near.sum())
    near_flip = int(np.sum((bs >= ec.CONF_THRESHOLD) != (cs >= ec.CONF_THRESHOLD)))

    # ---- end-to-end: instances after conf filter, matched by query index ----
    base_inst = ec.parse_instances(bl, bb, bs.astype(np.float32), num_classes=80)
    cand_inst = ec.parse_instances(cl, cb, cs.astype(np.float32), num_classes=80)
    base_q = {i["q"]: i for i in base_inst}
    cand_q = {i["q"]: i for i in cand_inst}
    shared_q = sorted(set(base_q) & set(cand_q))

    class_match = 0
    box_ious = []
    mask_ious = []
    OUT = 160  # compare masks in native 160-space to isolate model error from resize
    for q in shared_q:
        bi, ci = base_q[q], cand_q[q]
        if bi["classId"] == ci["classId"]:
            class_match += 1
        box_ious.append(ec.box_iou(bi["box"], ci["box"]))
        b_mask = bm[0, q] > ec.MASK_THRESHOLD
        c_mask = cm[0, q] > ec.MASK_THRESHOLD
        mask_ious.append(ec.mask_iou(b_mask, c_mask))

    return {
        "n_base_instances": len(base_inst),
        "n_cand_instances": len(cand_inst),
        "n_shared_queries": len(shared_q),
        "instance_count_delta": len(cand_inst) - len(base_inst),
        "label_agreement_all300": label_agree,
        "score_mae": score_mae,
        "score_max_abs_err": score_max,
        "box_mae": box_mae,
        "box_max_abs_err": box_max,
        "mask_logit_mae": mask_logit_mae,
        "mask_logit_max_abs_err": mask_logit_max,
        "mask_binary_flip_frac": mask_flip_frac,
        "near_conf_count": near_count,
        "near_conf_decision_flips": near_flip,
        "class_match_on_shared": class_match / len(shared_q) if shared_q else 1.0,
        "mean_box_iou_shared": float(np.mean(box_ious)) if box_ious else 1.0,
        "mean_mask_iou_shared": float(np.mean(mask_ious)) if mask_ious else 1.0,
        "min_mask_iou_shared": float(np.min(mask_ious)) if mask_ious else 1.0,
        "nan_or_inf": nan_inf,
    }


def aggregate(rows: list) -> dict:
    def m(key):
        return float(np.mean([r[key] for r in rows]))

    def mn(key):
        return float(np.min([r[key] for r in rows]))

    def mx(key):
        return float(np.max([r[key] for r in rows]))

    total_base = sum(r["n_base_instances"] for r in rows)
    total_cand = sum(r["n_cand_instances"] for r in rows)
    total_near = sum(r["near_conf_count"] for r in rows)
    total_near_flip = sum(r["near_conf_decision_flips"] for r in rows)
    return {
        "n_images": len(rows),
        "total_base_instances": total_base,
        "total_cand_instances": total_cand,
        "instance_recall_vs_base": (
            sum(r["n_shared_queries"] for r in rows) / total_base if total_base else 1.0
        ),
        "mean_label_agreement_all300": m("label_agreement_all300"),
        "mean_score_mae": m("score_mae"),
        "max_score_abs_err": mx("score_max_abs_err"),
        "mean_box_mae": m("box_mae"),
        "max_box_abs_err": mx("box_max_abs_err"),
        "mean_mask_logit_mae": m("mask_logit_mae"),
        "max_mask_logit_abs_err": mx("mask_logit_max_abs_err"),
        "mean_mask_binary_flip_frac": m("mask_binary_flip_frac"),
        "near_conf_total": total_near,
        "near_conf_decision_flips": total_near_flip,
        "mean_class_match_on_shared": m("class_match_on_shared"),
        "mean_box_iou_shared": m("mean_box_iou_shared"),
        "worst_box_iou_image_mean": mn("mean_box_iou_shared"),
        "mean_mask_iou_shared": m("mean_mask_iou_shared"),
        "worst_mask_iou_image_mean": mn("mean_mask_iou_shared"),
        "min_mask_iou_any_instance": mn("min_mask_iou_shared"),
        "any_nan_or_inf": any(r["nan_or_inf"] for r in rows),
    }


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--baseline", required=True)
    ap.add_argument("--candidate", required=True)
    ap.add_argument("--images", nargs="+", required=True)
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--json", dest="json_out", default=None)
    args = ap.parse_args()

    files = ec.list_images(args.images, limit=args.limit)
    if not files:
        raise SystemExit(f"No images under {args.images}")

    base = make_session(args.baseline)
    cand = make_session(args.candidate)

    rows = []
    for i, path in enumerate(files):
        x = ec.preprocess_file(path)
        row = compare_image(run(base, x), run(cand, x))
        row["image"] = os.path.basename(path)
        rows.append(row)
        if (i + 1) % 10 == 0:
            print(f"  {i+1}/{len(files)} images…", file=sys.stderr)

    summary = aggregate(rows)
    result = {
        "baseline": os.path.basename(args.baseline),
        "candidate": os.path.basename(args.candidate),
        "summary": summary,
        "per_image": rows,
    }
    print(json.dumps(summary, indent=2))
    if args.json_out:
        with open(args.json_out, "w") as fh:
            json.dump(result, fh, indent=2)
        print(f"\nwrote {args.json_out}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())