climperor/common.py

"""Shared helpers: config IO, slot geometry, crop preprocessing."""

import json
from pathlib import Path

import cv2
import numpy as np

ROOT = Path(__file__).parent
CONFIG_PATH = ROOT / "config.json"
TEMPLATES_REAL = ROOT / "templates" / "real"
TEMPLATES_CDN = ROOT / "templates" / "cdn"


def load_config() -> dict:
    with open(CONFIG_PATH, encoding="utf-8") as f:
        return json.load(f)


def save_config(cfg: dict) -> None:
    with open(CONFIG_PATH, "w", encoding="utf-8") as f:
        json.dump(cfg, f, ensure_ascii=False, indent=2)


def slot_rect_px(slot: dict, cfg: dict, img_w: int, img_h: int) -> tuple[int, int, int, int]:
    """Convert relative slot coords to pixel rect (x, y, w, h) for this image size."""
    w = cfg["slot_w_rel"] * img_h
    h = cfg["slot_h_rel"] * img_h
    cx = img_w / 2 + slot["cx_rel"] * img_h
    cy = slot["cy_rel"] * img_h
    return int(round(cx - w / 2)), int(round(cy - h / 2)), int(round(w)), int(round(h))


def crop_slot(img: np.ndarray, slot: dict, cfg: dict) -> np.ndarray | None:
    """Crop one slot, trim UI chrome (color bar / name plate), resize to canonical size."""
    ih, iw = img.shape[:2]
    x, y, w, h = slot_rect_px(slot, cfg, iw, ih)
    if w <= 0 or h <= 0:
        return None
    x, y = max(0, x), max(0, y)
    roi = img[y : min(y + h, ih), x : min(x + w, iw)]
    if roi.size == 0:
        return None

    t = cfg["crop_trim"]
    rh, rw = roi.shape[:2]
    y0 = int(rh * t["top"])
    y1 = int(rh * (1 - t["bottom"]))
    x0 = int(rw * t["left"])
    x1 = int(rw * (1 - t["right"]))
    inner = roi[y0:y1, x0:x1]
    if inner.size == 0:
        return None

    size = cfg["canonical_size"]
    return cv2.resize(inner, (size, size), interpolation=cv2.INTER_AREA)


def match_score(crop: np.ndarray, template: np.ndarray, mask: np.ndarray | None = None) -> float:
    """Normalized cross-correlation between two same-sized BGR images.

    If mask is given (uint8, nonzero = use), only those pixels contribute.
    Used to ignore the ranked-medal banner that sits on the bottom/right of
    every top-bar portrait in ranked matchmaking.
    """
    if crop.shape != template.shape:
        template = cv2.resize(template, (crop.shape[1], crop.shape[0]), interpolation=cv2.INTER_AREA)
    if mask is None:
        res = cv2.matchTemplate(crop, template, cv2.TM_CCOEFF_NORMED)
        return float(res[0][0])

    if mask.shape[:2] != crop.shape[:2]:
        mask = cv2.resize(mask, (crop.shape[1], crop.shape[0]), interpolation=cv2.INTER_NEAREST)
    sel = mask > 0
    if int(sel.sum()) < 32:
        return -1.0
    a = crop[sel].astype(np.float32).ravel()
    b = template[sel].astype(np.float32).ravel()
    a -= a.mean()
    b -= b.mean()
    denom = float(np.linalg.norm(a) * np.linalg.norm(b))
    return float(a @ b / denom) if denom > 1e-6 else -1.0


def ranked_match_mask(size: int, cfg: dict) -> np.ndarray:
    """Canonical-size mask that zeroes the bottom rank bar and right medal."""
    rm = cfg.get("match", {}).get("ranked_mask", {})
    bottom = float(rm.get("bottom", 0.32))
    right = float(rm.get("right", 0.22))
    mask = np.ones((size, size), np.uint8) * 255
    mask[int(size * (1.0 - bottom)) :, :] = 0
    mask[:, int(size * (1.0 - right)) :] = 0
    return mask


def has_ranked_overlay(img: np.ndarray, cfg: dict) -> bool:
    """True when most slots show the gold rank medal on the right edge.

    Bot / unranked strategy-time frames have no medals, so this stays false
    and recognition keeps using the full portrait.
    """
    if not cfg.get("slots"):
        return False
    ih, iw = img.shape[:2]
    hits = 0
    checked = 0
    for slot in cfg["slots"]:
        x, y, w, h = slot_rect_px(slot, cfg, iw, ih)
        if w <= 0 or h <= 0:
            continue
        roi = img[max(0, y) : min(ih, y + h), max(0, x) : min(iw, x + w)]
        if roi.size == 0:
            continue
        checked += 1
        rh, rw = roi.shape[:2]
        corner = roi[int(rh * 0.35) :, int(rw * 0.68) :]
        if corner.size == 0:
            continue
        hsv = cv2.cvtColor(corner, cv2.COLOR_BGR2HSV)
        gold = cv2.inRange(hsv, (8, 70, 90), (40, 255, 255))
        if float(gold.mean()) > 18.0:
            hits += 1
    return checked > 0 and hits >= max(6, checked * 0.6)


def load_template_library() -> list[tuple[str, str, np.ndarray]]:
    """Return list of (hero_key, source, image). source is 'real' or 'cdn'."""
    lib: list[tuple[str, str, np.ndarray]] = []
    if TEMPLATES_REAL.is_dir():
        for hero_dir in sorted(TEMPLATES_REAL.iterdir()):
            if not hero_dir.is_dir():
                continue
            for png in hero_dir.glob("*.png"):
                img = cv2.imread(str(png))
                if img is not None:
                    lib.append((hero_dir.name, "real", img))
    if TEMPLATES_CDN.is_dir():
        for png in TEMPLATES_CDN.glob("*.png"):
            img = cv2.imread(str(png))
            if img is not None:
                lib.append((png.stem, "cdn", img))
    return lib