vosonandCursor
f32d24b8f8
从 dota2-draft-vision 迁出并定名，作为天梯选将识别项目起点。 Co-authored-by: Cursor <cursoragent@cursor.com>
120 lines
3.8 KiB
Python
120 lines
3.8 KiB
Python
"""Read the game-mode label under the top-center timer (e.g. 全英雄选择)."""
|
|
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from common import ROOT
|
|
|
|
TEMPLATES = ROOT / "templates" / "modes"
|
|
|
|
# key -> Chinese label as drawn under the timer during hero selection
|
|
MODES = {
|
|
"all_pick": "全英雄选择",
|
|
"captains_mode": "队长模式",
|
|
"random_draft": "随机征召",
|
|
"single_draft": "单一征召",
|
|
"ability_draft": "技能征召",
|
|
}
|
|
|
|
|
|
def mode_roi(img: np.ndarray, cfg: dict) -> np.ndarray:
|
|
"""Crop the strip under the draft timer where the mode name sits."""
|
|
m = cfg.get("mode_label", {})
|
|
ih, iw = img.shape[:2]
|
|
y0 = int(ih * m.get("y0_rel", 0.045))
|
|
y1 = int(ih * m.get("y1_rel", 0.072))
|
|
x0 = int(iw * m.get("x0_rel", 0.40))
|
|
x1 = int(iw * m.get("x1_rel", 0.60))
|
|
return img[y0:y1, x0:x1]
|
|
|
|
|
|
def _ink(roi: np.ndarray) -> np.ndarray:
|
|
"""Binary mask of the bright mode glyphs on the dark header."""
|
|
if roi.size == 0:
|
|
return np.zeros((1, 1), np.uint8)
|
|
gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
|
|
return (gray > 160).astype(np.uint8) * 255
|
|
|
|
|
|
def _tight(mask: np.ndarray, height: int = 28) -> np.ndarray | None:
|
|
ys, xs = np.where(mask > 0)
|
|
if len(xs) < 8:
|
|
return None
|
|
crop = mask[ys.min():ys.max() + 1, xs.min():xs.max() + 1]
|
|
h, w = crop.shape
|
|
nh = height
|
|
nw = max(8, int(round(w * (nh / h))))
|
|
return cv2.resize(crop, (nw, nh), interpolation=cv2.INTER_AREA)
|
|
|
|
|
|
def load_mode_templates() -> dict[str, np.ndarray]:
|
|
out = {}
|
|
if not TEMPLATES.is_dir():
|
|
return out
|
|
for p in TEMPLATES.glob("*.png"):
|
|
img = cv2.imread(str(p), cv2.IMREAD_GRAYSCALE)
|
|
if img is not None:
|
|
out[p.stem] = img
|
|
return out
|
|
|
|
|
|
def detect_mode(img: np.ndarray, cfg: dict, templates: dict | None = None) -> dict | None:
|
|
"""Return {key, label, score} or None."""
|
|
templates = templates if templates is not None else load_mode_templates()
|
|
if not templates:
|
|
return None
|
|
ink = _tight(_ink(mode_roi(img, cfg)))
|
|
if ink is None:
|
|
return None
|
|
best_key, best = None, -1.0
|
|
for key, tmpl in templates.items():
|
|
h = min(ink.shape[0], tmpl.shape[0])
|
|
a = cv2.resize(ink, (max(8, int(ink.shape[1] * h / ink.shape[0])), h))
|
|
b = cv2.resize(tmpl, (max(8, int(tmpl.shape[1] * h / tmpl.shape[0])), h))
|
|
big, small = (a, b) if a.shape[1] >= b.shape[1] else (b, a)
|
|
if big.shape[0] < small.shape[0] or big.shape[1] < small.shape[1]:
|
|
continue
|
|
score = float(cv2.matchTemplate(big, small, cv2.TM_CCOEFF_NORMED).max())
|
|
if score > best:
|
|
best, best_key = score, key
|
|
min_score = cfg.get("mode_label", {}).get("min_score", 0.55)
|
|
if best_key is None or best < min_score:
|
|
return None
|
|
return {"key": best_key, "label": MODES.get(best_key, best_key), "score": round(best, 3)}
|
|
|
|
|
|
def build_template(img: np.ndarray, cfg: dict, key: str) -> Path:
|
|
"""Save a mode template from a live selection frame."""
|
|
TEMPLATES.mkdir(parents=True, exist_ok=True)
|
|
ink = _tight(_ink(mode_roi(img, cfg)))
|
|
if ink is None:
|
|
raise SystemExit("no mode glyphs found in ROI - check mode_label coords")
|
|
out = TEMPLATES / f"{key}.png"
|
|
cv2.imwrite(str(out), ink)
|
|
return out
|
|
|
|
|
|
def _main() -> None:
|
|
import sys
|
|
|
|
from common import load_config
|
|
|
|
cfg = load_config()
|
|
if len(sys.argv) < 2:
|
|
raise SystemExit("usage: python modes.py <frame.png> [--build KEY]")
|
|
img = cv2.imread(sys.argv[1])
|
|
if img is None:
|
|
raise SystemExit(f"cannot read {sys.argv[1]}")
|
|
if "--build" in sys.argv:
|
|
key = sys.argv[sys.argv.index("--build") + 1]
|
|
print(build_template(img, cfg, key))
|
|
return
|
|
found = detect_mode(img, cfg)
|
|
print(found or "no mode matched")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
_main()
|