FEnigma/tools/eval_map_vision.py

#!/usr/bin/env python3
"""Evaluate map_vision against the hand annotations, and render overlays.

Usage:
  .venv/bin/python tools/eval_map_vision.py                 # score the fixtures
  .venv/bin/python tools/eval_map_vision.py shot.png [...]  # try your own images

Two scores are reported, because the lenient one hid a real error:

  cell    fraction of annotated points landing in the correct cell. Too
          forgiving on its own: the annotations sit near cell centres, so a
          grid wrong by a whole line still passes.
  spread  every annotated point should sit at (col + a, row + b) in recovered
          grid coordinates for ONE constant (a, b) -- the label's offset
          inside its cell. So fit that constant and report the spread of the
          residual, in cell units. A skewed or off-by-one-line grid shows up
          here even when every point is nominally in the right cell.
"""
import json
import sys
from pathlib import Path

import cv2
import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from fenigma import map_vision as mv          # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
SHOTS = ROOT / "tests" / "fixtures" / "map_shots"
GT_PATH = ROOT / "tests" / "fixtures" / "map_shots_gt.json"


def draw(img, sol, markers):
    vis = img.copy()
    h, w = vis.shape[:2]
    q = np.linalg.inv(sol.H) @ np.array([[0, w, w, 0], [0, 0, h, h], [1, 1, 1, 1]], np.float64)
    ij = q[:2] / q[2]
    for i in range(int(np.floor(ij[0].min())) - 1, int(np.ceil(ij[0].max())) + 2):
        for j in range(int(np.floor(ij[1].min())) - 1, int(np.ceil(ij[1].max())) + 2):
            col, row = sol.si * i + sol.du, sol.sj * j + sol.dv
            if not (0 <= col < mv.COLS and 1 <= row <= mv.ROWS):
                continue
            p = sol.H @ np.array([[i, i + 1, i + 1, i], [j, j, j + 1, j + 1], [1, 1, 1, 1.0]])
            if np.any(np.abs(p[2]) < 1e-9):
                continue
            p = (p[:2] / p[2]).T
            cv2.polylines(vis, [p.astype(np.int32)], True, (0, 255, 255), 2, cv2.LINE_AA)
            lx = i + (mv.PAD_L if sol.si > 0 else 1 - mv.PAD_L)
            ly = j + (mv.PAD_T if sol.sj > 0 else 1 - mv.PAD_T)
            t = sol.H @ np.array([lx, ly, 1.0])
            if abs(t[2]) < 1e-9:
                continue
            x, y = int(t[0] / t[2]), int(t[1] / t[2])
            f = max(0.45, min(2.0, float(np.linalg.norm(p[1] - p[0])) / 190.0))
            txt = f"{mv.LARGE_X[col]}{row}"
            cv2.putText(vis, txt, (x, y), cv2.FONT_HERSHEY_SIMPLEX, f, (0, 0, 0), int(f * 5) + 2)
            cv2.putText(vis, txt, (x, y), cv2.FONT_HERSHEY_SIMPLEX, f, (0, 255, 0), int(f * 2) + 1)
    for m in markers:
        x, y, bw, bh = m["box"]
        c = (0, 0, 255) if m["side"] == "hostile" else (255, 210, 0)
        cv2.rectangle(vis, (x, y), (x + bw, y + bh), c, 2)
        txt = m["coord"] + (f' {m["unit"]}' if m.get("unit") else "")
        cv2.putText(vis, txt, (x, y - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 4)
        cv2.putText(vis, txt, (x, y - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.55, c, 2)
    return vis


def run_one(path, outdir):
    """Solve a single arbitrary screenshot (no ground truth needed)."""
    sol, img, err = mv.solve_path(path)
    if sol is None:
        print(f"{path.name}: REJECTED - {err}")
        return 1
    markers = mv.find_markers(img, sol)
    print(f"{path.name}: solved, {sol.votes} label votes, "
          f"cell {sol.steps[0]:.0f}x{sol.steps[1]:.0f}px, {len(markers)} markers")
    for m in markers:
        print(f"   {m['side']:<8} {m['coord']:<9} {m.get('unit') or 'unknown type'}")
    out = outdir / f"solved_{path.stem}.png"
    cv2.imwrite(str(out), draw(img, sol, markers))
    print(f"   overlay: {out}")
    return 0


def main():
    args = [a for a in sys.argv[1:]]
    outdir = ROOT / "build" / "map_vision"
    images = [Path(a) for a in args if Path(a).suffix.lower() in (".png", ".jpg", ".jpeg")]
    if images:
        outdir.mkdir(parents=True, exist_ok=True)
        rc = 0
        for p in images:
            rc |= run_one(p, outdir)
        return rc
    outdir = Path(args[0]) if args else outdir
    outdir.mkdir(parents=True, exist_ok=True)
    gt = json.loads(GT_PATH.read_text())
    names = sorted(k for k in gt if k.endswith(".png"))
    tiles, ok_t, n_t = [], 0, 0
    print(f"{'shot':<8} {'pts':>4} {'cell':>6} {'spread':>7} {'votes':>5} "
          f"{'mk':>3}  status")
    for name in names:
        path = SHOTS / name
        native_w = cv2.imread(str(path), cv2.IMREAD_REDUCED_COLOR_8).shape[1] * 8
        pts = gt[name]
        n_t += len(pts)
        sol, img, err = mv.solve_path(path)
        s = img.shape[1] / native_w
        if sol is None:
            print(f"{name:<8} {len(pts):>4} {'-':>6} {'-':>7} {'-':>5} {'-':>3}  REJECT: {err}")
            vis = img.copy()
            cv2.putText(vis, f"{name} REJECTED", (12, 34), cv2.FONT_HERSHEY_SIMPLEX,
                        1.0, (0, 0, 255), 3)
            cv2.putText(vis, err[:56], (12, 66), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
            tiles.append(vis)
            cv2.imwrite(str(outdir / f"rejected_{name}"), vis)
            continue
        offs, right = [], 0
        for lab, x, y in pts:
            got = sol.cell_of(x * s, y * s)
            if got and got[0] == lab:
                right += 1
            g = sol.grid_of(x * s, y * s)
            if g is None:
                continue
            offs.append((g[0] - mv.LARGE_X.index(lab[0]), g[1] - int(lab[1:])))
        ok_t += right
        o = np.array(offs)
        spread = float(np.sqrt(((o - o.mean(axis=0)) ** 2).sum(axis=1).mean()))
        markers = mv.find_markers(img, sol)
        verdict = "GOOD" if spread < 0.08 else ("SKEWED" if spread < 0.3 else "BAD")
        print(f"{name:<8} {len(pts):>4} {right/len(pts):>5.0%} {spread:>7.3f} "
              f"{sol.votes:>5} {len(markers):>3}  {verdict}")
        vis = draw(img, sol, markers)
        cv2.putText(vis, f"{name}  {len(markers)} markers  spread {spread:.3f}",
                    (12, 34), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 3)
        tiles.append(vis)
        cv2.imwrite(str(outdir / f"solved_{name}"), vis)
        for m in markers:
            print(f"           {m['side']:<8} {m['coord']:<9} "
                  f"{str(m.get('unit')):<26} s={m['unit_score']:.2f} d={m['unit_margin']:.3f}")
    print(f"\nTOTAL {ok_t}/{n_t} annotated points in the correct cell "
          f"({ok_t / max(n_t, 1):.0%})")
    TW, TH, cols = 860, 520, 3
    rows = (len(tiles) + cols - 1) // cols
    sheet = np.zeros((TH * rows, TW * cols, 3), np.uint8)
    for i, t in enumerate(tiles):
        hh, ww = t.shape[:2]
        sc = min(TW / ww, TH / hh)
        t2 = cv2.resize(t, (int(ww * sc), int(hh * sc)))
        y, x = (i // cols) * TH, (i % cols) * TW
        sheet[y:y + t2.shape[0], x:x + t2.shape[1]] = t2
    cv2.imwrite(str(outdir / "sheet.png"), sheet)
    print(f"wrote {outdir / 'sheet.png'}")


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