166 lines
4.0 KiB
Python
166 lines
4.0 KiB
Python
import os
|
||
import cv2
|
||
import numpy as np
|
||
|
||
# =========================================================
|
||
# 全局配置
|
||
# =========================================================
|
||
|
||
TEMPLATE_DIR = "./muban_image1"
|
||
|
||
# 三个 ROI（x, y, w, h）
|
||
ROI_1 = (782, 614, 164, 128)
|
||
ROI_2 = (837, 791, 100, 99)
|
||
ROI_3 = (873, 736, 141, 110)
|
||
|
||
USE_GRAY = False
|
||
TM_METHOD = cv2.TM_CCOEFF_NORMED
|
||
SCORE_THRESH = 0.4
|
||
|
||
|
||
# =========================================================
|
||
# 模板加载
|
||
# =========================================================
|
||
def load_templates(template_dir):
|
||
templates = []
|
||
if not os.path.isdir(template_dir):
|
||
raise RuntimeError(f"模板目录不存在: {template_dir}")
|
||
|
||
for fname in os.listdir(template_dir):
|
||
if not fname.lower().endswith((".png", ".jpg", ".jpeg", ".bmp")):
|
||
continue
|
||
|
||
path = os.path.join(template_dir, fname)
|
||
img = cv2.imread(path)
|
||
if img is None:
|
||
continue
|
||
|
||
if USE_GRAY:
|
||
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
|
||
h, w = img.shape[:2]
|
||
name = os.path.splitext(fname)[0]
|
||
|
||
templates.append({
|
||
"name": name,
|
||
"img": img,
|
||
"h": h,
|
||
"w": w
|
||
})
|
||
|
||
print(f"[INFO] 已加载 {len(templates)} 个模板")
|
||
return templates
|
||
|
||
|
||
TEMPLATES = load_templates(TEMPLATE_DIR)
|
||
|
||
|
||
# =========================================================
|
||
# ROI 裁剪
|
||
# =========================================================
|
||
def crop_roi(img, roi):
|
||
x, y, w, h = roi
|
||
return img[y:y + h, x:x + w]
|
||
|
||
|
||
# =========================================================
|
||
# 单 ROI 模板匹配
|
||
# =========================================================
|
||
def match_one_roi(img, roi):
|
||
roi_img = crop_roi(img, roi)
|
||
if roi_img.size == 0:
|
||
return None, 0.0
|
||
|
||
if USE_GRAY:
|
||
roi_img = cv2.cvtColor(roi_img, cv2.COLOR_BGR2GRAY)
|
||
|
||
best_name = None
|
||
best_score = -1.0
|
||
|
||
for t in TEMPLATES:
|
||
if t["h"] > roi_img.shape[0] or t["w"] > roi_img.shape[1]:
|
||
continue
|
||
|
||
res = cv2.matchTemplate(roi_img, t["img"], TM_METHOD)
|
||
_, max_val, _, _ = cv2.minMaxLoc(res)
|
||
|
||
if max_val > best_score:
|
||
best_score = max_val
|
||
best_name = t["name"]
|
||
|
||
if best_score < SCORE_THRESH:
|
||
return None, best_score
|
||
|
||
return best_name, best_score
|
||
|
||
|
||
# =========================================================
|
||
# 核心逻辑：ROI_1 → ROI_2 → ROI_3
|
||
# =========================================================
|
||
def match_image(img):
|
||
roi_list = [
|
||
("roi1", ROI_1),
|
||
("roi2", ROI_2),
|
||
("roi3", ROI_3),
|
||
]
|
||
|
||
for roi_name, roi in roi_list:
|
||
name, score = match_one_roi(img, roi)
|
||
if name is not None:
|
||
return {
|
||
"result": name,
|
||
"roi": roi_name,
|
||
"score": score
|
||
}
|
||
|
||
return {
|
||
"result": None,
|
||
"roi": None,
|
||
"score": 0.0
|
||
}
|
||
|
||
|
||
# =========================================================
|
||
# 可视化（只画命中的 ROI）
|
||
# =========================================================
|
||
def draw_result(img, result):
|
||
vis = img.copy()
|
||
|
||
roi_map = {
|
||
"roi1": ROI_1,
|
||
"roi2": ROI_2,
|
||
"roi3": ROI_3
|
||
}
|
||
|
||
if result["roi"] not in roi_map:
|
||
return vis
|
||
|
||
x, y, w, h = roi_map[result["roi"]]
|
||
text = f'{result["result"]} {result["score"]:.2f}'
|
||
|
||
cv2.rectangle(vis, (x, y), (x + w, y + h), (0, 255, 0), 2)
|
||
cv2.putText(
|
||
vis, text, (x, y - 10),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.7, (0, 255, 0), 2
|
||
)
|
||
|
||
return vis
|
||
|
||
|
||
# =========================================================
|
||
# main 测试
|
||
# =========================================================
|
||
if __name__ == "__main__":
|
||
img = cv2.imread("1.png")
|
||
if img is None:
|
||
raise RuntimeError("测试图片未找到")
|
||
|
||
result = match_image(img)
|
||
print("[RESULT]", result)
|
||
|
||
vis = draw_result(img, result)
|
||
cv2.imshow("template match", vis)
|
||
cv2.waitKey(0)
|
||
cv2.destroyAllWindows()
|