climperor/roles.py

"""Read role-queue lane labels under top-bar portraits.

Row under each portrait (role-queue only, own team): 优势路 / 中路 / 劣势路 /
辅助 / 纯辅助. Your own slot index comes from GSI team_slot, not from name tint.

The role text is flat grey with zero saturation, so a threshold on
value+saturation isolates it cleanly. Matching is done on the binary mask
(icon included) rather than OCR: there are only five possible strings and
they differ in width, so mask IoU separates them by a wide margin.
"""

import cv2
import numpy as np

from common import ROOT, slot_rect_px

TEMPLATES_ROLES = ROOT / "templates" / "roles"

# key -> (in-game text, lane position number)
ROLES = {
    "safe": ("优势路", 1),
    "mid": ("中路", 2),
    "off": ("劣势路", 3),
    "soft_support": ("辅助", 4),
    "hard_support": ("纯辅助", 5),
}

# canonical mask geometry, chosen so 1440p text (~17px tall) upsamples slightly
STRIP_H = 24
STRIP_W = 160


def _row_rect(slot: dict, cfg: dict, img_w: int, img_h: int, row: str) -> tuple[int, int, int, int]:
    r = cfg["text_rows"][row]
    x, _, w, _ = slot_rect_px(slot, cfg, img_w, img_h)
    pad = int(round(w * 0.35))  # names/roles overflow the portrait width
    y0 = int(round(r["y0_rel"] * img_h))
    y1 = int(round(r["y1_rel"] * img_h))
    return x - pad, y0, w + 2 * pad, y1 - y0


def _text_mask(patch: np.ndarray, min_value: int, max_sat: float) -> np.ndarray:
    """Isolate the flat light-grey glyphs from the dark blurred background."""
    p = patch.astype(np.float32)
    mx = p.max(axis=2)
    mn = p.min(axis=2)
    sat = (mx - mn) / np.maximum(mx, 1.0)
    return ((mx > min_value) & (sat < max_sat)).astype(np.uint8) * 255


def _tight(mask: np.ndarray) -> np.ndarray | None:
    """Crop to the ink, then normalize height so resolution stops mattering."""
    ys, xs = np.nonzero(mask)
    if ys.size < 40:
        return None
    m = mask[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1]
    h, w = m.shape
    scale = STRIP_H / h
    m = cv2.resize(m, (max(1, int(round(w * scale))), STRIP_H), interpolation=cv2.INTER_AREA)
    canvas = np.zeros((STRIP_H, STRIP_W), np.uint8)
    m = m[:, :STRIP_W]
    canvas[:, : m.shape[1]] = m
    return (canvas > 127).astype(np.uint8) * 255


def role_mask(img: np.ndarray, slot: dict, cfg: dict) -> np.ndarray | None:
    """Binary mask of one slot's role label, or None when there is no label."""
    ih, iw = img.shape[:2]
    x, y, w, h = _row_rect(slot, cfg, iw, ih, "role")
    patch = img[max(0, y) : min(ih, y + h), max(0, x) : min(iw, x + w)]
    if patch.size == 0:
        return None
    t = cfg["text_rows"]["role"]
    return _tight(_text_mask(patch, t.get("min_value", 110), t.get("max_sat", 0.08)))


def name_tint(img: np.ndarray, slot: dict, cfg: dict) -> tuple[float, float] | None:
    """Mean value and saturation of the name glyphs: (value, saturation)."""
    ih, iw = img.shape[:2]
    x, y, w, h = _row_rect(slot, cfg, iw, ih, "name")
    patch = img[max(0, y) : min(ih, y + h), max(0, x) : min(iw, x + w)]
    if patch.size == 0:
        return None
    p = patch.astype(np.float32)
    mx = p.max(axis=2)
    thr = max(90.0, float(mx.max()) * 0.7)
    sel = mx > thr
    if int(sel.sum()) < 30:
        return None
    px = p[sel]
    hi = px.max(axis=1)
    lo = px.min(axis=1)
    return float(hi.mean()), float(((hi - lo) / np.maximum(hi, 1.0)).mean())


def iou(a: np.ndarray, b: np.ndarray) -> float:
    ab = a > 0
    bb = b > 0
    union = int((ab | bb).sum())
    return float((ab & bb).sum()) / union if union else 0.0


def load_role_templates() -> dict[str, np.ndarray]:
    if not TEMPLATES_ROLES.is_dir():
        return {}
    out = {}
    for key in ROLES:
        f = TEMPLATES_ROLES / f"{key}.png"
        if f.is_file():
            img = cv2.imread(str(f), cv2.IMREAD_GRAYSCALE)
            if img is not None:
                out[key] = img
    return out


def detect_roles(img: np.ndarray, cfg: dict, templates: dict[str, np.ndarray] | None = None) -> dict:
    """Per-slot role labels from the role-queue text under top-bar portraits.

    Returns {"self_team": str|None, "roles": {slot: {...}}}.
    Your own top-bar slot comes from GSI team_slot elsewhere - this helper
    does not guess it from name brightness.
    Slots without a role label (the enemy team, or any non-role-queue mode)
    are simply absent from "roles".
    """
    if templates is None:
        templates = load_role_templates()
    cutoff = cfg.get("roles", {}).get("min_iou", 0.55)

    roles: dict[int, dict] = {}
    for slot in cfg.get("slots", []):
        mask = role_mask(img, slot, cfg)
        if mask is None:
            continue
        ranked = sorted(((iou(mask, t), k) for k, t in templates.items()), reverse=True)
        if not ranked or ranked[0][0] < cutoff:
            continue
        score, key = ranked[0]
        roles[slot["index"]] = {
            "role": key,
            "label": ROLES[key][0],
            "position": ROLES[key][1],
            "score": round(score, 3),
        }

    self_team = None
    if roles:
        self_team = "radiant" if min(roles) <= 5 else "dire"

    return {"self_team": self_team, "roles": roles}


def _main() -> None:
    """python roles.py <frame.png>                 - report roles found
    python roles.py <frame.png> --build off,safe,mid,soft_support,hard_support
                                                   - save templates from slots 1..N
    """
    import sys

    from common import load_config

    args = sys.argv[1:]
    if not args:
        print(_main.__doc__)
        return
    frame = cv2.imread(args[0])
    if frame is None:
        raise SystemExit(f"cannot read {args[0]}")
    cfg = load_config()

    if "--build" in args:
        labels = args[args.index("--build") + 1].split(",")
        TEMPLATES_ROLES.mkdir(parents=True, exist_ok=True)
        for slot, key in zip(cfg["slots"], labels):
            key = key.strip()
            if key not in ROLES:
                raise SystemExit(f"unknown role {key!r}, expected one of {list(ROLES)}")
            mask = role_mask(frame, slot, cfg)
            if mask is None:
                raise SystemExit(f"slot {slot['index']} has no role text")
            out = TEMPLATES_ROLES / f"{key}.png"
            cv2.imwrite(str(out), mask)
            print(f"slot {slot['index']} -> {key} ({ROLES[key][0]})  {out}")
        return

    import json

    print(json.dumps(detect_roles(frame, cfg), ensure_ascii=False, indent=1))


if __name__ == "__main__":
    _main()