94 lines
2.7 KiB
Python
94 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MatchResult:
|
|
confidence: float
|
|
bbox: tuple[int, int, int, int]
|
|
passed: bool
|
|
method: str
|
|
debug: dict[str, Any]
|
|
|
|
|
|
def _as_gray(image: np.ndarray) -> np.ndarray:
|
|
if image.ndim == 2:
|
|
return image
|
|
return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
|
|
|
|
def match_template(
|
|
image: np.ndarray,
|
|
template: np.ndarray,
|
|
*,
|
|
threshold: float = 0.85,
|
|
method: int = cv2.TM_CCOEFF_NORMED,
|
|
) -> MatchResult:
|
|
if image.size == 0:
|
|
raise ValueError("image is empty")
|
|
if template.size == 0:
|
|
raise ValueError("template is empty")
|
|
if template.shape[0] > image.shape[0] or template.shape[1] > image.shape[1]:
|
|
raise ValueError("template cannot be larger than image")
|
|
|
|
image_gray = _as_gray(image)
|
|
template_gray = _as_gray(template)
|
|
response = cv2.matchTemplate(image_gray, template_gray, method)
|
|
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(response)
|
|
|
|
if method in (cv2.TM_SQDIFF, cv2.TM_SQDIFF_NORMED):
|
|
top_left = min_loc
|
|
confidence = 1.0 - float(min_val)
|
|
else:
|
|
top_left = max_loc
|
|
confidence = float(max_val)
|
|
|
|
width = int(template.shape[1])
|
|
height = int(template.shape[0])
|
|
bbox = (int(top_left[0]), int(top_left[1]), width, height)
|
|
return MatchResult(
|
|
confidence=confidence,
|
|
bbox=bbox,
|
|
passed=confidence >= threshold,
|
|
method=_method_name(method),
|
|
debug={
|
|
"threshold": threshold,
|
|
"min_value": float(min_val),
|
|
"max_value": float(max_val),
|
|
"min_location": tuple(map(int, min_loc)),
|
|
"max_location": tuple(map(int, max_loc)),
|
|
"image_shape": tuple(map(int, image.shape)),
|
|
"template_shape": tuple(map(int, template.shape)),
|
|
},
|
|
)
|
|
|
|
|
|
def save_match_debug(image: np.ndarray, result: MatchResult, output_path: str | Path) -> Path:
|
|
output = Path(output_path)
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
marked = image.copy()
|
|
x, y, width, height = result.bbox
|
|
color = (0, 255, 0) if result.passed else (0, 0, 255)
|
|
cv2.rectangle(marked, (x, y), (x + width, y + height), color, 2)
|
|
cv2.imwrite(str(output), marked)
|
|
return output
|
|
|
|
|
|
def _method_name(method: int) -> str:
|
|
names = {
|
|
cv2.TM_CCOEFF: "TM_CCOEFF",
|
|
cv2.TM_CCOEFF_NORMED: "TM_CCOEFF_NORMED",
|
|
cv2.TM_CCORR: "TM_CCORR",
|
|
cv2.TM_CCORR_NORMED: "TM_CCORR_NORMED",
|
|
cv2.TM_SQDIFF: "TM_SQDIFF",
|
|
cv2.TM_SQDIFF_NORMED: "TM_SQDIFF_NORMED",
|
|
}
|
|
return names.get(method, str(method))
|