climperor/pc/autocalibrate.py

"""Locate the 10 top-bar hero slots automatically, no manual box drawing.

The top bar puts a player-coloured strip above every portrait, and those ten
colours are fixed by the game. Finding them gives both the horizontal position
and the slot order for free, at any resolution.

Usage:
    python autocalibrate.py samples/raw/draft_141704.png
    python autocalibrate.py samples/raw/draft_141704.png --check   # inspect only

Writes slot geometry into config.json and preview/autocalibrate_check.png.
"""

import sys
from pathlib import Path

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

import time

import cv2
import numpy as np

from common import ROOT, load_config, save_config

PREVIEW_DIR = ROOT / "preview"

# Dota 2 player colours as RGB: radiant slots 1-5 then dire slots 6-10
PLAYER_COLORS = [
    (51, 117, 255),
    (102, 255, 191),
    (191, 0, 191),
    (243, 240, 11),
    (255, 107, 0),
    (254, 134, 194),
    (161, 180, 71),
    (101, 217, 247),
    (0, 131, 33),
    (164, 105, 0),
]

COLOR_TOLERANCE = 60


def find_color_bar_rows(img: np.ndarray) -> tuple[int, int]:
    """Rows spanned by the player-colour strips."""
    h, w = img.shape[:2]
    targets = np.array([(b, g, r) for (r, g, b) in PLAYER_COLORS], dtype=np.int16)
    search = img[: int(h * 0.08)].astype(np.int16)

    hits = []
    for y in range(search.shape[0]):
        d = np.linalg.norm(search[y][:, None, :] - targets[None, :, :], axis=2)
        hits.append(int((d.min(axis=1) < COLOR_TOLERANCE).sum()))

    hits = np.array(hits)
    strong = np.where(hits > w * 0.15)[0]
    if strong.size == 0:
        raise SystemExit(
            "no player colour bars found - is this really a draft/strategy-time frame?"
        )
    runs = np.split(strong, np.where(np.diff(strong) > 2)[0] + 1)
    run = max(runs, key=len)
    return int(run[0]), int(run[-1])


def find_slots(img: np.ndarray, y0: int, y1: int) -> tuple[list[float], float]:
    """Slot centre x for all ten slots, plus the common slot width."""
    targets = np.array([(b, g, r) for (r, g, b) in PLAYER_COLORS], dtype=np.int16)
    band = np.median(img[y0 : y1 + 1].astype(np.int16), axis=0)

    d = np.linalg.norm(band[:, None, :] - targets[None, :, :], axis=2)
    best, dist = d.argmin(axis=1), d.min(axis=1)
    ok = dist < COLOR_TOLERANCE

    centers: list[float | None] = []
    widths: list[int | None] = []
    for idx in range(10):
        xs = np.where(ok & (best == idx))[0]
        if xs.size == 0:
            centers.append(None)
            widths.append(None)
            continue
        runs = np.split(xs, np.where(np.diff(xs) > 5)[0] + 1)
        run = max(runs, key=len)
        centers.append(float(run[0] + run[-1]) / 2)
        widths.append(int(run[-1] - run[0] + 1))

    if sum(c is not None for c in centers) < 8:
        raise SystemExit("found fewer than 8 colour bars - frame is probably not a full top bar")

    # A bar whose colour bleeds into the portrait behind it comes out too wide,
    # and its centre is then wrong by several pixels. Least squares would let
    # such a bar drag the whole row; judge each bar by its width first and only
    # trust the well-formed ones.
    width = float(np.median([wd for wd in widths if wd is not None]))
    reliable = [
        c is not None and wd is not None and abs(wd - width) <= width * 0.15
        for c, wd in zip(centers, widths)
    ]

    # slot pitch is identical for both teams, so take it from every good pair
    steps = [
        (centers[j] - centers[i]) / (j - i)
        for team in (range(0, 5), range(5, 10))
        for i in team
        for j in team
        if j > i and reliable[i] and reliable[j]
    ]
    if not steps:
        raise SystemExit("no reliable colour bars to measure slot spacing from")
    pitch = float(np.median(steps))

    fitted: list[float] = []
    for team in (range(0, 5), range(5, 10)):
        idx = [i for i in team if reliable[i]] or list(team)
        base = float(np.median([centers[i] - pitch * (i - team[0]) for i in idx]))
        fitted += [base + pitch * (i - team[0]) for i in team]

    return fitted, width


def find_portrait_bottom(img: np.ndarray, bar_bottom: int, centers: list[float], width: float) -> int:
    """Row where the portraits give way to the name plates.

    Uses the brightness gap between portrait columns and the gaps between
    portraits: it is large while portraits are present and collapses to zero
    the moment they end. A plain row-to-row delta does not work here because
    the player names further down produce an even bigger jump.
    """
    h, w = img.shape[:2]
    half = width / 2

    inside = np.concatenate(
        [np.arange(int(c - half) + 6, int(c + half) - 6) for c in centers]
    )
    gaps = np.concatenate(
        [
            np.arange(int(a + half) + 10, int(b - half) - 10)
            for team in (centers[:5], centers[5:])
            for a, b in zip(team[:-1], team[1:])
        ]
    )
    inside = inside[(inside >= 0) & (inside < w)]
    gaps = gaps[(gaps >= 0) & (gaps < w)]

    top = bar_bottom + 1
    end = min(h, bar_bottom + int(h * 0.15))
    strip = img[top:end].astype(np.int16)
    contrast = np.abs(strip[:, inside].mean(axis=(1, 2)) - strip[:, gaps].mean(axis=(1, 2)))

    faded = np.where(contrast < contrast.max() * 0.05)[0]
    if faded.size == 0:
        raise SystemExit("could not find the bottom edge of the portraits")
    return top + int(faded[0])


def main() -> None:
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    path = sys.argv[1]
    check_only = "--check" in sys.argv

    img = cv2.imread(path)
    if img is None:
        sys.exit(f"cannot read image: {path}")
    h, w = img.shape[:2]

    bar_top, bar_bottom = find_color_bar_rows(img)
    centers, width = find_slots(img, bar_top, bar_bottom)
    portrait_top = bar_bottom + 1
    portrait_bottom = find_portrait_bottom(img, bar_bottom, centers, width)
    height = portrait_bottom - portrait_top

    print(f"image           : {w}x{h}")
    print(f"colour bar rows : {bar_top}-{bar_bottom}")
    print(f"portrait rows   : {portrait_top}-{portrait_bottom}  (height {height})")
    print(f"slot width      : {width:.0f}")
    print(f"slot centres    : {', '.join(f'{c:.0f}' for c in centers)}")

    if height < 20 or width < 20:
        sys.exit("detected geometry looks wrong - refusing to write config")

    cy = portrait_top + height / 2
    slots = [
        {"index": i + 1, "cx_rel": (c - w / 2) / h, "cy_rel": cy / h}
        for i, c in enumerate(centers)
    ]

    PREVIEW_DIR.mkdir(exist_ok=True)
    check = img.copy()
    for s, c in zip(slots, centers):
        x0, x1 = int(c - width / 2), int(c + width / 2)
        cv2.rectangle(check, (x0, portrait_top), (x1, portrait_bottom), (0, 0, 255), 2)
        cv2.putText(check, str(s["index"]), (x0 + 4, portrait_bottom + 26),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    cv2.imwrite(str(PREVIEW_DIR / "autocalibrate_check.png"), check[: portrait_bottom + 40])

    tiles = [
        cv2.copyMakeBorder(
            img[portrait_top:portrait_bottom, int(c - width / 2) : int(c + width / 2)],
            2, 2, 2, 2, cv2.BORDER_CONSTANT, value=(0, 0, 255),
        )
        for c in centers
    ]
    cv2.imwrite(str(PREVIEW_DIR / "autocalibrate_slots.png"), np.hstack(tiles))
    print("wrote preview/autocalibrate_check.png and preview/autocalibrate_slots.png")

    if check_only:
        print("--check given, config.json untouched")
        return

    cfg = load_config()
    cfg["calibrated_on"] = time.strftime("%Y-%m-%d %H:%M:%S")
    cfg["calibrated_from"] = str(Path(path).name)
    cfg["slots"] = slots
    cfg["slot_w_rel"] = width / h
    cfg["slot_h_rel"] = height / h
    # the ROI is already just the portrait, so nothing left to trim away
    cfg["crop_trim"] = {"top": 0.0, "bottom": 0.0, "left": 0.0, "right": 0.0}
    save_config(cfg)
    print("config.json updated")


if __name__ == "__main__":
    main()