188 lines
6.9 KiB
Python
188 lines
6.9 KiB
Python
"""模板识别 + 标定工具测试（用合成图片，无需真机）。需要 opencv/numpy。"""
|
||
|
||
import numpy as np
|
||
import pytest
|
||
|
||
cv2 = pytest.importorskip("cv2")
|
||
|
||
from hbc.board import CellKind # noqa: E402
|
||
from hbc.calibrate import _fit_scale, _map_box_to_original # noqa: E402
|
||
from hbc.calibrate import main as calib_main # noqa: E402
|
||
from hbc.config import ( # noqa: E402
|
||
GridConfig,
|
||
TemplateMatchConfig,
|
||
level_layout_path,
|
||
load_level_grid,
|
||
)
|
||
from hbc.recognizer import recognize_board # noqa: E402
|
||
from hbc.templates import TemplateLibrary # noqa: E402
|
||
|
||
CELL = 40
|
||
|
||
# 合成棋盘布局：数字=正常块颜色, 'W'=墙(石), 'I'=墙(冰), '.'=空
|
||
LAYOUT = [
|
||
["0", "1", "2", "W"],
|
||
["1", "I", "0", "."],
|
||
["2", "0", "1", "2"],
|
||
]
|
||
|
||
# 每种外观对应的 BGR 纯色块（足够区分即可）
|
||
APPEARANCE = {
|
||
"0": (60, 60, 220), # 红
|
||
"1": (60, 200, 60), # 绿
|
||
"2": (220, 120, 60), # 蓝
|
||
"W": (90, 90, 90), # 石墙(灰)
|
||
"I": (230, 220, 180), # 冰墙(浅蓝白)
|
||
".": (10, 10, 10), # 空(近黑)
|
||
}
|
||
|
||
|
||
def _swatch(bgr):
|
||
img = np.zeros((CELL, CELL, 3), np.uint8)
|
||
img[:] = bgr
|
||
return img
|
||
|
||
|
||
def _build_board_image():
|
||
rows, cols = len(LAYOUT), len(LAYOUT[0])
|
||
img = np.zeros((rows * CELL, cols * CELL, 3), np.uint8)
|
||
for r in range(rows):
|
||
for c in range(cols):
|
||
img[r * CELL:(r + 1) * CELL, c * CELL:(c + 1) * CELL] = _swatch(APPEARANCE[LAYOUT[r][c]])
|
||
return img
|
||
|
||
|
||
def _write_templates(root):
|
||
(root / "normal").mkdir(parents=True)
|
||
(root / "wall").mkdir(parents=True)
|
||
(root / "empty").mkdir(parents=True)
|
||
cv2.imwrite(str(root / "normal" / "red.png"), _swatch(APPEARANCE["0"]))
|
||
cv2.imwrite(str(root / "normal" / "green.png"), _swatch(APPEARANCE["1"]))
|
||
cv2.imwrite(str(root / "normal" / "blue.png"), _swatch(APPEARANCE["2"]))
|
||
cv2.imwrite(str(root / "wall" / "stone.png"), _swatch(APPEARANCE["W"]))
|
||
cv2.imwrite(str(root / "wall" / "ice.png"), _swatch(APPEARANCE["I"]))
|
||
cv2.imwrite(str(root / "empty" / "hole.png"), _swatch(APPEARANCE["."]))
|
||
|
||
|
||
def test_template_recognition_multi_wall(tmp_path):
|
||
root = tmp_path / "templates"
|
||
_write_templates(root)
|
||
lib = TemplateLibrary.load(TemplateMatchConfig(templates_dir=str(root), canonical_size=32))
|
||
|
||
img = _build_board_image()
|
||
rows, cols = len(LAYOUT), len(LAYOUT[0])
|
||
grid = GridConfig(x=0, y=0, w=cols * CELL, h=rows * CELL, rows=rows, cols=cols)
|
||
board = recognize_board(img, grid, lib)
|
||
|
||
# 校验每格类型
|
||
for r in range(rows):
|
||
for c in range(cols):
|
||
cell = board.get(r, c)
|
||
token = LAYOUT[r][c]
|
||
if token in ("0", "1", "2"):
|
||
assert cell.kind == CellKind.NORMAL, (r, c, cell)
|
||
elif token in ("W", "I"):
|
||
assert cell.kind == CellKind.BRICK, (r, c, cell)
|
||
else:
|
||
assert cell.kind == CellKind.EMPTY, (r, c, cell)
|
||
|
||
# 两种墙应被区分到不同 meta['wall']
|
||
assert board.get(0, 3).meta["wall"] == "stone"
|
||
assert board.get(1, 1).meta["wall"] == "ice"
|
||
|
||
# 相同外观的正常块应得到相同 color id（求解器据此判等）
|
||
assert board.color_at(0, 0) == board.color_at(1, 2) # 两个红
|
||
assert board.color_at(0, 1) == board.color_at(1, 0) # 两个绿
|
||
assert board.color_at(0, 0) != board.color_at(0, 1) # 红 != 绿
|
||
|
||
|
||
def test_low_confidence_becomes_ice_wall(tmp_path):
|
||
"""置信度不足（如冰块/未登记方块）一律判为冰墙，而非空格。"""
|
||
root = tmp_path / "templates"
|
||
_write_templates(root) # 只有红/绿/蓝 + 墙 + 空
|
||
lib = TemplateLibrary.load(TemplateMatchConfig(templates_dir=str(root), canonical_size=32))
|
||
|
||
purple = _swatch((200, 40, 200)) # 库里没有的外观 -> 低置信度
|
||
cell = lib.classify(purple)
|
||
assert cell.is_brick and cell.meta.get("wall") == "ice"
|
||
|
||
# 真正的空块（与空模板一致，高分）仍判为空格
|
||
empty_cell = lib.classify(_swatch(APPEARANCE["."]))
|
||
assert empty_cell.is_empty and not empty_cell.is_brick
|
||
|
||
|
||
def test_layout_save_load(tmp_path):
|
||
p = tmp_path / "layout.json"
|
||
g = GridConfig(x=12, y=34, w=400, h=300, rows=6, cols=8)
|
||
g.save_json(p)
|
||
g2 = GridConfig.load_json(p)
|
||
assert g2.to_dict() == g.to_dict()
|
||
|
||
|
||
def test_level_layout_files(tmp_path):
|
||
"""7 关网格逐渐变大：每关一个 level{N}.json，文件名带关号，统一放 layouts/。"""
|
||
d = tmp_path / "layouts"
|
||
GridConfig(0, 0, 200, 200, 5, 5).save_json(_mk(level_layout_path(1, d)))
|
||
GridConfig(0, 0, 240, 240, 6, 6).save_json(level_layout_path(2, d))
|
||
GridConfig(0, 0, 280, 280, 8, 8).save_json(level_layout_path(7, d))
|
||
|
||
assert level_layout_path(1, d).name == "level1.json"
|
||
assert load_level_grid(1, d).rows == 5
|
||
assert load_level_grid(2, d).cols == 6
|
||
assert load_level_grid(7, d).w == 280
|
||
# 未标定的关卡应报错
|
||
import pytest as _pt
|
||
with _pt.raises(FileNotFoundError):
|
||
load_level_grid(3, d)
|
||
|
||
|
||
def _mk(p):
|
||
p.parent.mkdir(parents=True, exist_ok=True)
|
||
return p
|
||
|
||
|
||
def test_roi_scale_mapping():
|
||
"""缩放窗口里选的框，要能正确换算回原图坐标。"""
|
||
# 原图 1800 高，限制 900 -> scale=0.5
|
||
scale = _fit_scale(1000, 1800, 900)
|
||
assert abs(scale - 0.5) < 1e-9
|
||
# 在缩放图上选 (50,100,200,300) -> 原图应为两倍
|
||
assert _map_box_to_original((50, 100, 200, 300), scale) == (100, 200, 400, 600)
|
||
# 不需要缩放时原样返回
|
||
assert _fit_scale(400, 300, 900) == 1.0
|
||
assert _map_box_to_original((1, 2, 3, 4), 1.0) == (1, 2, 3, 4)
|
||
|
||
|
||
def test_calibrate_export_and_preview(tmp_path):
|
||
img_path = tmp_path / "board.png"
|
||
cv2.imwrite(str(img_path), _build_board_image())
|
||
layouts_dir = tmp_path / "layouts"
|
||
rows, cols = len(LAYOUT), len(LAYOUT[0])
|
||
GridConfig(0, 0, cols * CELL, rows * CELL, rows, cols).save_json(_mk(level_layout_path(1, layouts_dir)))
|
||
|
||
cells_dir = tmp_path / "cells"
|
||
rc = calib_main(["export", "--image", str(img_path), "--level", "1",
|
||
"--layouts-dir", str(layouts_dir), "--out", str(cells_dir)])
|
||
assert rc == 0
|
||
exported = list(cells_dir.glob("*.png"))
|
||
assert len(exported) == rows * cols
|
||
|
||
preview = tmp_path / "preview.png"
|
||
rc = calib_main(["preview", "--image", str(img_path), "--level", "1",
|
||
"--layouts-dir", str(layouts_dir), "--out", str(preview)])
|
||
assert rc == 0
|
||
assert preview.exists()
|
||
|
||
|
||
def test_calibrate_roi_with_box(tmp_path):
|
||
"""无 GUI：用 --box 直接给坐标，写入指定关卡文件。"""
|
||
img_path = tmp_path / "board.png"
|
||
cv2.imwrite(str(img_path), _build_board_image())
|
||
layouts_dir = tmp_path / "layouts"
|
||
rc = calib_main(["roi", "--image", str(img_path), "--level", "3",
|
||
"--rows", "5", "--cols", "5", "--box", "0,0,200,200",
|
||
"--layouts-dir", str(layouts_dir)])
|
||
assert rc == 0
|
||
g = load_level_grid(3, layouts_dir)
|
||
assert (g.x, g.y, g.w, g.h, g.rows, g.cols) == (0, 0, 200, 200, 5, 5)
|