vosonandCursor
f32d24b8f8
从 dota2-draft-vision 迁出并定名，作为天梯选将识别项目起点。 Co-authored-by: Cursor <cursoragent@cursor.com>
79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
"""Crop slots from a labeled screenshot and add them to the real template library.
|
|
|
|
Usage:
|
|
# 1) preview: crop 10 slots to preview/ so you can see what each slot contains
|
|
python build_library.py samples/shot1.png
|
|
|
|
# 2) import: provide 10 comma-separated hero keys (left to right), '?' to skip a slot
|
|
python build_library.py samples/shot1.png tinker,earthshaker,juggernaut,dazzle,vengefulspirit,axe,sniper,slark,lion,drow_ranger
|
|
|
|
Hero keys must match Steam internal names without the npc_dota_hero_ prefix
|
|
(see heroes.json after running fetch_cdn_templates.py).
|
|
"""
|
|
|
|
import json
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
|
|
from common import ROOT, TEMPLATES_REAL, crop_slot, load_config
|
|
|
|
PREVIEW_DIR = ROOT / "preview"
|
|
|
|
|
|
def known_hero_keys() -> set[str]:
|
|
path = ROOT / "heroes.json"
|
|
if not path.exists():
|
|
return set()
|
|
with open(path, encoding="utf-8") as f:
|
|
return {h["key"] for h in json.load(f)}
|
|
|
|
|
|
def main() -> None:
|
|
if len(sys.argv) < 2:
|
|
sys.exit(__doc__)
|
|
image_path = sys.argv[1]
|
|
labels = sys.argv[2].split(",") if len(sys.argv) > 2 else None
|
|
|
|
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")
|
|
|
|
crops = [(slot["index"], crop_slot(img, slot, cfg)) for slot in cfg["slots"]]
|
|
|
|
if labels is None:
|
|
PREVIEW_DIR.mkdir(exist_ok=True)
|
|
for idx, crop in crops:
|
|
if crop is not None:
|
|
cv2.imwrite(str(PREVIEW_DIR / f"slot_{idx}.png"), crop)
|
|
print(f"wrote {len(crops)} crops to {PREVIEW_DIR}/ - inspect them, then rerun with labels")
|
|
return
|
|
|
|
if len(labels) != 10:
|
|
sys.exit(f"expected 10 labels, got {len(labels)}")
|
|
known = known_hero_keys()
|
|
stamp = time.strftime("%Y%m%d_%H%M%S")
|
|
added = 0
|
|
for (idx, crop), label in zip(crops, labels):
|
|
label = label.strip()
|
|
if label == "?" or crop is None:
|
|
continue
|
|
if known and label not in known:
|
|
print(f" WARNING slot {idx}: '{label}' not in heroes.json - saved anyway, double-check spelling")
|
|
hero_dir = TEMPLATES_REAL / label
|
|
hero_dir.mkdir(parents=True, exist_ok=True)
|
|
cv2.imwrite(str(hero_dir / f"{stamp}_s{idx}.png"), crop)
|
|
added += 1
|
|
total = sum(1 for _ in TEMPLATES_REAL.rglob("*.png"))
|
|
heroes = sum(1 for d in TEMPLATES_REAL.iterdir() if d.is_dir())
|
|
print(f"added {added} templates; library now {total} images / {heroes} heroes")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|