climperor/pc/recognize.py

"""Recognize the 10 drafted heroes from a strategy-time screenshot.

Usage:
    python recognize.py samples/raw/draft.png
    python recognize.py samples/raw/draft.png --truth tinker,earthshaker,...,drow_ranger
    python recognize.py samples/raw/draft.png --sheet

Outputs per-slot JSON with top-1 hero, score and margin; slots failing the
confidence gate are reported as null. With --truth, prints accuracy and saves
misrecognized crops to failures/ (debug only, gitignored).

recognize_image() is the reusable entry point used by gsi_watch.py.
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import json
import time

import cv2
import numpy as np

from common import (
    ROOT,
    crop_slot,
    has_ranked_overlay,
    load_config,
    load_template_library,
    match_score,
    ranked_match_mask,
)

FAILURES_DIR = ROOT / "failures"
PREVIEW_DIR = ROOT / "preview"


def write_sheet(img: np.ndarray, results: list[dict], cfg: dict) -> str:
    """Contact sheet of every slot with its predicted hero, for eyeballing."""
    scale = 2
    tiles = []
    for r in results:
        crop = crop_slot(img, cfg["slots"][r["slot"] - 1], cfg)
        if crop is None:
            continue
        tile = cv2.resize(crop, None, fx=scale, fy=scale, interpolation=cv2.INTER_LANCZOS4)
        label = np.zeros((54, tile.shape[1], 3), np.uint8)
        name = r["hero"] or f"?{r['raw_best']}"
        colour = (120, 255, 120) if r["hero"] else (120, 200, 255)
        cv2.putText(label, f"{r['slot']} {name[:16]}", (4, 20),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.42, colour, 1, cv2.LINE_AA)
        cv2.putText(label, f"s{r['score']:.2f} m{r['margin']:.2f}", (4, 42),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.42, (170, 170, 170), 1, cv2.LINE_AA)
        stack = np.vstack([tile, label])
        tiles.append(cv2.copyMakeBorder(stack, 2, 2, 2, 2, cv2.BORDER_CONSTANT, value=(60, 60, 60)))

    PREVIEW_DIR.mkdir(exist_ok=True)
    out = PREVIEW_DIR / "recognize_sheet.png"
    cv2.imwrite(str(out), np.hstack(tiles))
    return str(out)


def recognize_slot(crop, library, cfg, mask=None):
    """Return (best_hero, best_score, margin, scored list).

    cfg is accepted for call-site compatibility; score gates are applied by the caller.
    """
    _ = cfg
    best_per_hero: dict[str, float] = {}
    for hero, tmpl in library:
        s = match_score(crop, tmpl, mask)
        if s > best_per_hero.get(hero, -2.0):
            best_per_hero[hero] = s
    ranked = sorted(best_per_hero.items(), key=lambda kv: kv[1], reverse=True)
    if not ranked:
        return None, 0.0, 0.0, []
    top1 = ranked[0]
    margin = top1[1] - ranked[1][1] if len(ranked) > 1 else 1.0
    return top1[0], top1[1], margin, ranked[:3]


def recognize_image(img: np.ndarray, cfg: dict | None = None, library=None) -> dict:
    """Recognize all slots in a full-screen frame.

    cfg and library are accepted so a long-running caller can load the
    template library once instead of on every frame.

    Ranked matchmaking draws a title bar + medal over every portrait; when
    that overlay is detected we match only the unoccluded face region.
    """
    cfg = cfg if cfg is not None else load_config()
    library = library if library is not None else load_template_library()

    t0 = time.perf_counter()
    min_score = cfg["match"]["min_score"]
    min_margin = cfg["match"]["min_margin"]
    ranked_ui = has_ranked_overlay(img, cfg)
    mask = ranked_match_mask(cfg["canonical_size"], cfg) if ranked_ui else None

    results = []
    for slot in cfg["slots"]:
        crop = crop_slot(img, slot, cfg)
        if crop is None:
            results.append({"slot": slot["index"], "hero": None, "score": 0, "margin": 0, "top3": []})
            continue
        hero, score, margin, top3 = recognize_slot(crop, library, cfg, mask)
        passed = score >= min_score and margin >= min_margin
        results.append(
            {
                "slot": slot["index"],
                "hero": hero if passed else None,
                "raw_best": hero,
                "score": round(score, 3),
                "margin": round(margin, 3),
                "top3": [[h, round(s, 3)] for h, s in top3],
            }
        )
    elapsed = time.perf_counter() - t0

    return {
        "radiant": results[:5],
        "dire": results[5:],
        "slots": results,
        "recognized": sum(1 for r in results if r["hero"]),
        "ranked_overlay": ranked_ui,
        "library_size": len(library),
        "elapsed_ms": round(elapsed * 1000),
    }


def main() -> None:
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    image_path = sys.argv[1]
    truth = None
    if "--truth" in sys.argv:
        truth = sys.argv[sys.argv.index("--truth") + 1].split(",")
        if len(truth) != 10:
            sys.exit(f"--truth expects 10 comma-separated keys, got {len(truth)}")

    img = cv2.imread(image_path)
    if img is None:
        sys.exit(f"cannot read image: {image_path}")
    cfg = load_config()
    if not cfg["slots"]:
        sys.exit("config.json has no slots - run calibrate.py first")
    library = load_template_library()
    if not library:
        sys.exit("template library is empty - run fetch_cdn_templates.py")

    out = recognize_image(img, cfg, library)
    results = out.pop("slots")
    print(json.dumps(out, ensure_ascii=False, indent=1))

    if "--sheet" in sys.argv:
        print(f"sheet: {write_sheet(img, results, cfg)}")

    if truth:
        FAILURES_DIR.mkdir(exist_ok=True)
        stamp = time.strftime("%Y%m%d_%H%M%S")
        correct = 0
        for r, expected in zip(results, truth):
            expected = expected.strip()
            got = r["hero"]
            ok = got == expected
            correct += ok
            mark = "OK " if ok else "ERR"
            print(f"{mark} slot {r['slot']}: expected={expected} got={got} (raw={r.get('raw_best')} score={r['score']} margin={r['margin']})")
            if not ok:
                slot_cfg = cfg["slots"][r["slot"] - 1]
                crop = crop_slot(img, slot_cfg, cfg)
                if crop is not None:
                    cv2.imwrite(str(FAILURES_DIR / f"{stamp}_s{r['slot']}_{expected}.png"), crop)
        print(f"accuracy: {correct}/10, misses saved to failures/ (filename contains the correct key)")


if __name__ == "__main__":
    main()