338 lines
12 KiB
Python
338 lines
12 KiB
Python
"""GPU-accelerated template matching using FFT correlation + integral images.
|
||
|
||
Math is identical to OpenCV's TM_SQDIFF_NORMED + TM_CCOEFF_NORMED, but the
|
||
heavy operations are restructured to exploit GPU strengths:
|
||
|
||
* Cross-correlation SI = sum(I[y+u, x+v] * T[u, v]) is computed in the
|
||
frequency domain. We FFT the drawing ONCE, then per-variant pad+FFT the
|
||
template and elementwise-multiply. On Blackwell (sm_120), cuDNN's direct
|
||
conv kernels are unoptimized and take ~1 s for an 8000² drawing with a
|
||
100² kernel — FFT correlation does the same work in ~12 ms.
|
||
|
||
* Local sums S1 = sum(I in window) and S2 = sum(I² in window) are computed
|
||
from integral images via four-corner subtraction — O(1) per output pixel,
|
||
so per-variant cost is essentially free.
|
||
|
||
Putting these together: per-variant cost drops from ~3 s (direct conv) to
|
||
~15 ms, and the cost is dominated by the one-time drawing FFT (~230 ms on
|
||
an 8000×6000 page).
|
||
|
||
API and return shape match the CPU implementation in counting.py.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import Iterable
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import torch
|
||
import torch.fft as tfft
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
DEFAULT_THRESHOLD = 0.7
|
||
MAX_DRAWING_PX = 9000
|
||
EPS = 1e-6
|
||
|
||
# Binarization: pixels darker than this become "ink". Kept at 245 — measured
|
||
# the actual renders: drawing strokes are crisp (well below 245) so tightening
|
||
# barely changes the drawing ink (2.10% vs 2.24% at 210 vs 245), but faint
|
||
# templates (e.g. a low-contrast symbol crop, 241–255) lose ALL ink below ~245
|
||
# and silently match nothing. Not worth the fragility — the false-positive fix
|
||
# is the consensus scoring + verify gate below, not binarization.
|
||
INK_THRESHOLD = 245
|
||
|
||
# Propose-then-verify tuning.
|
||
# - Consensus uses min(sim_sq, cc): far stricter than the old max(), so we
|
||
# propose slightly below the user threshold to protect recall …
|
||
PROPOSE_RELAX = 0.08
|
||
# - … then a geometric ink-overlap gate removes the candidates that only
|
||
# matched by luck (right ink density, wrong shape).
|
||
VERIFY_INK_IOU = 0.25
|
||
|
||
|
||
def _prep(img_path: Path) -> np.ndarray:
|
||
arr = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE)
|
||
if arr is None:
|
||
raise RuntimeError(f"Could not load {img_path}")
|
||
_, ink = cv2.threshold(arr, INK_THRESHOLD, 255, cv2.THRESH_BINARY_INV)
|
||
return ink
|
||
|
||
|
||
def _crop_to_content(template: np.ndarray) -> np.ndarray:
|
||
mask = template > 0
|
||
if not mask.any():
|
||
return template
|
||
ys, xs = np.where(mask)
|
||
y0, y1 = ys.min(), ys.max() + 1
|
||
x0, x1 = xs.min(), xs.max() + 1
|
||
pad = 2
|
||
y0 = max(0, y0 - pad)
|
||
x0 = max(0, x0 - pad)
|
||
y1 = min(template.shape[0], y1 + pad)
|
||
x1 = min(template.shape[1], x1 + pad)
|
||
return template[y0:y1, x0:x1]
|
||
|
||
|
||
def _rotate(img: np.ndarray, angle_deg: int) -> np.ndarray:
|
||
if angle_deg == 0:
|
||
return img
|
||
if angle_deg in (90, 180, 270):
|
||
return np.rot90(img, k=angle_deg // 90)
|
||
h, w = img.shape[:2]
|
||
center = (w / 2, h / 2)
|
||
M = cv2.getRotationMatrix2D(center, angle_deg, 1.0)
|
||
cos = abs(M[0, 0]); sin = abs(M[0, 1])
|
||
new_w = int(h * sin + w * cos)
|
||
new_h = int(h * cos + w * sin)
|
||
M[0, 2] += new_w / 2 - center[0]
|
||
M[1, 2] += new_h / 2 - center[1]
|
||
return cv2.warpAffine(img, M, (new_w, new_h),
|
||
flags=cv2.INTER_NEAREST, borderValue=0)
|
||
|
||
|
||
def _nms(boxes: list[tuple], scores: list[float],
|
||
overlap_thresh: float = 0.3) -> list[int]:
|
||
if not boxes:
|
||
return []
|
||
boxes_arr = np.array(boxes, dtype=np.float32)
|
||
x1 = boxes_arr[:, 0]; y1 = boxes_arr[:, 1]
|
||
x2 = boxes_arr[:, 2]; y2 = boxes_arr[:, 3]
|
||
areas = (x2 - x1) * (y2 - y1)
|
||
order = np.argsort(scores)[::-1]
|
||
keep = []
|
||
while order.size > 0:
|
||
i = order[0]
|
||
keep.append(int(i))
|
||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||
w = np.maximum(0.0, xx2 - xx1)
|
||
h = np.maximum(0.0, yy2 - yy1)
|
||
inter = w * h
|
||
union = areas[i] + areas[order[1:]] - inter
|
||
iou = inter / np.maximum(union, EPS)
|
||
order = order[1:][iou <= overlap_thresh]
|
||
return keep
|
||
|
||
|
||
def _integral(I: torch.Tensor) -> torch.Tensor:
|
||
"""Build integral image. Output shape: (H+1, W+1), with zero first row/col.
|
||
|
||
II[i, j] = sum(I[:i, :j]). Then for any rectangle [y, y+h) × [x, x+w):
|
||
local_sum = II[y+h, x+w] - II[y, x+w] - II[y+h, x] + II[y, x]
|
||
"""
|
||
H, W = I.shape
|
||
II = torch.zeros((H + 1, W + 1), device=I.device, dtype=I.dtype)
|
||
II[1:, 1:] = I.cumsum(0).cumsum(1)
|
||
return II
|
||
|
||
|
||
def _local_sum(II: torch.Tensor, h: int, w: int) -> torch.Tensor:
|
||
"""Sum of every h×w window using integral image. Output: (H-h+1, W-w+1)."""
|
||
return (II[h:, w:] - II[:-h, w:] - II[h:, :-w] + II[:-h, :-w])
|
||
|
||
|
||
def _fft_correlate(
|
||
Ifft: torch.Tensor, T: torch.Tensor, H: int, W: int,
|
||
) -> torch.Tensor:
|
||
"""Cross-correlation of drawing (whose rfft2 is Ifft) with template T.
|
||
|
||
Returns the 'valid' region [0, H-h+1) × [0, W-w+1) — matching the
|
||
output region of cv2.matchTemplate / F.conv2d.
|
||
|
||
Math: cross-correlation in time domain = FFT(I) * conj(FFT(T_padded))
|
||
in frequency domain. T_padded has T in the top-left of a zero canvas
|
||
of size (H, W).
|
||
"""
|
||
h, w = T.shape
|
||
Tpad = torch.zeros((H, W), device=T.device, dtype=T.dtype)
|
||
Tpad[:h, :w] = T
|
||
Tfft = tfft.rfft2(Tpad)
|
||
full = tfft.irfft2(Ifft * Tfft.conj(), s=(H, W))
|
||
return full[: H - h + 1, : W - w + 1]
|
||
|
||
|
||
def count_template(
|
||
template_path: Path,
|
||
drawing_path: Path,
|
||
threshold: float = DEFAULT_THRESHOLD,
|
||
rotations: Iterable[int] = (0, 45, 90, 135, 180, 225, 270, 315),
|
||
scales: Iterable[float] = (0.6, 0.75, 0.9, 1.0, 1.15, 1.35, 1.6),
|
||
exclude_box: tuple | None = None,
|
||
area_tolerance: float = 200.0, # ignored — kept for API compat
|
||
mirror: bool = True,
|
||
) -> dict:
|
||
"""GPU template matcher. Same return shape as the CPU implementation."""
|
||
device = torch.device("cuda")
|
||
|
||
template = _prep(template_path)
|
||
template = _crop_to_content(template)
|
||
drawing = _prep(drawing_path)
|
||
|
||
coord_scale = 1.0
|
||
if max(drawing.shape) > MAX_DRAWING_PX:
|
||
coord_scale = max(drawing.shape) / MAX_DRAWING_PX
|
||
new_w = int(drawing.shape[1] / coord_scale)
|
||
new_h = int(drawing.shape[0] / coord_scale)
|
||
drawing = cv2.resize(drawing, (new_w, new_h),
|
||
interpolation=cv2.INTER_AREA)
|
||
|
||
if exclude_box is not None:
|
||
ex_x, ex_y, ex_w, ex_h = exclude_box
|
||
ex_x = int(ex_x / coord_scale); ex_y = int(ex_y / coord_scale)
|
||
ex_w = int(ex_w / coord_scale); ex_h = int(ex_h / coord_scale)
|
||
H_, W_ = drawing.shape
|
||
ex_x = max(0, min(W_, ex_x)); ex_y = max(0, min(H_, ex_y))
|
||
ex_x2 = max(0, min(W_, ex_x + ex_w))
|
||
ex_y2 = max(0, min(H_, ex_y + ex_h))
|
||
drawing[ex_y:ex_y2, ex_x:ex_x2] = 0
|
||
|
||
if min(template.shape) < 8 or int((template > 0).sum()) < 5:
|
||
return {"count": 0, "matches": [], "threshold_used": threshold}
|
||
|
||
H, W = drawing.shape
|
||
|
||
# One-time GPU setup. Use float32 normalized to [0,1] to keep magnitudes
|
||
# tame in the FFT (large summations on uint8 → 0..255 can overflow into
|
||
# noisier float regions of the spectrum).
|
||
drawing_t = torch.from_numpy(drawing.astype(np.float32) / 255.0).to(device)
|
||
drawing_sq_t = drawing_t * drawing_t
|
||
|
||
# FFT of drawing (used by every variant for SI)
|
||
Ifft = tfft.rfft2(drawing_t)
|
||
|
||
# Integral images for S1 and S2 (used by every variant)
|
||
II = _integral(drawing_t)
|
||
II2 = _integral(drawing_sq_t)
|
||
|
||
# Build all variants
|
||
base_variants = [template]
|
||
if mirror:
|
||
base_variants.append(cv2.flip(template, 1))
|
||
|
||
variants_np: list[np.ndarray] = []
|
||
for variant in base_variants:
|
||
for rot in rotations:
|
||
rotated = _rotate(variant, rot)
|
||
for scale in scales:
|
||
nw = max(8, int(rotated.shape[1] * scale))
|
||
nh = max(8, int(rotated.shape[0] * scale))
|
||
if nh > H or nw > W:
|
||
continue
|
||
tmpl = cv2.resize(rotated, (nw, nh),
|
||
interpolation=cv2.INTER_AREA)
|
||
variants_np.append(tmpl)
|
||
|
||
if not variants_np:
|
||
return {"count": 0, "matches": [], "threshold_used": threshold}
|
||
|
||
all_boxes: list[tuple] = []
|
||
all_scores: list[float] = []
|
||
all_vidx: list[int] = [] # which variant produced each candidate (verify)
|
||
|
||
propose_thr = max(0.0, threshold - PROPOSE_RELAX)
|
||
|
||
for vi, v in enumerate(variants_np):
|
||
v_norm = (v.astype(np.float32) / 255.0)
|
||
h_t, w_t = v_norm.shape
|
||
T = torch.from_numpy(v_norm).to(device)
|
||
|
||
# SI = sum_{u,v} I[y+u, x+v] * T[u, v]
|
||
SI = _fft_correlate(Ifft, T, H, W)
|
||
|
||
# S1 = sum_{u,v} I[y+u, x+v]
|
||
# S2 = sum_{u,v} I[y+u, x+v]^2
|
||
S1 = _local_sum(II, h_t, w_t)
|
||
S2 = _local_sum(II2, h_t, w_t)
|
||
|
||
n = float(h_t * w_t)
|
||
sumT = float(T.sum())
|
||
sumT2 = float((T * T).sum())
|
||
meanT = sumT / n
|
||
var_T_sum = sumT2 - n * meanT * meanT
|
||
|
||
# TM_SQDIFF_NORMED → similarity = 1 - normalized squared difference
|
||
sq = (sumT2 - 2.0 * SI + S2) / torch.sqrt(
|
||
torch.clamp(sumT2 * S2, min=EPS))
|
||
sim_sq = 1.0 - sq
|
||
|
||
if var_T_sum <= EPS:
|
||
sim = sim_sq
|
||
else:
|
||
# TM_CCOEFF_NORMED
|
||
cc_num = SI - meanT * S1
|
||
var_I_sum = torch.clamp(S2 - S1 * S1 / n, min=0.0)
|
||
cc = cc_num / torch.sqrt(
|
||
torch.clamp(var_T_sum * var_I_sum, min=EPS))
|
||
# Score = CCOEFF_NORMED only. The old code used max(sim_sq, cc),
|
||
# which let sim_sq's spurious highs over big white/low-ink regions
|
||
# through → thousands of false positives. cc is the shape-
|
||
# discriminative metric and keeps recall; the ink-IoU verify gate
|
||
# below removes cc's remaining thin-line spurious hits. (min() was
|
||
# tried and over-pruned: sim_sq is structurally low for sparse
|
||
# linework even on true matches, so it killed real symbols too.)
|
||
sim = cc
|
||
|
||
hits = sim >= propose_thr
|
||
if hits.any():
|
||
ys, xs = torch.where(hits)
|
||
scores_t = sim[ys, xs]
|
||
ys = ys.cpu().numpy(); xs = xs.cpu().numpy()
|
||
scores_np = scores_t.cpu().numpy()
|
||
for y, x, s in zip(ys, xs, scores_np):
|
||
all_boxes.append(
|
||
(float(x), float(y), float(x + w_t), float(y + h_t)))
|
||
all_scores.append(float(s))
|
||
all_vidx.append(vi)
|
||
|
||
del SI, S1, S2, sq, sim_sq, T
|
||
|
||
del drawing_t, drawing_sq_t, Ifft, II, II2
|
||
torch.cuda.empty_cache()
|
||
|
||
if not all_boxes:
|
||
return {"count": 0, "matches": [], "threshold_used": threshold}
|
||
|
||
keep = _nms(all_boxes, all_scores, overlap_thresh=0.3)
|
||
|
||
# ── Verify pass ──────────────────────────────────────────────
|
||
# For each surviving box: (a) enforce the real user threshold (propose was
|
||
# relaxed for recall), and (b) require the candidate window's ink to
|
||
# geometrically overlap the matched template's ink. (b) is what removes the
|
||
# leftover "right ink density, wrong shape" false positives.
|
||
matches = []
|
||
for i in keep:
|
||
score = all_scores[i]
|
||
if score < threshold:
|
||
continue
|
||
x0, y0, x1, y1 = all_boxes[i]
|
||
tmpl = variants_np[all_vidx[i]]
|
||
th, tw = tmpl.shape
|
||
yi, xi = int(round(y0)), int(round(x0))
|
||
if yi < 0 or xi < 0 or yi + th > H or xi + tw > W:
|
||
continue
|
||
win = drawing[yi:yi + th, xi:xi + tw]
|
||
if win.shape != tmpl.shape:
|
||
continue
|
||
t_ink = tmpl > 0
|
||
w_ink = win > 0
|
||
union = int(np.logical_or(t_ink, w_ink).sum())
|
||
if union == 0:
|
||
continue
|
||
iou = int(np.logical_and(t_ink, w_ink).sum()) / union
|
||
if iou < VERIFY_INK_IOU:
|
||
continue
|
||
matches.append({
|
||
"x": int(x0 * coord_scale),
|
||
"y": int(y0 * coord_scale),
|
||
"w": int((x1 - x0) * coord_scale),
|
||
"h": int((y1 - y0) * coord_scale),
|
||
"score": round(score, 3),
|
||
})
|
||
return {"count": len(matches), "matches": matches,
|
||
"threshold_used": threshold}
|