基于FFT的快速模板匹配：从数学证明到高效实现

=== 数值对比（与 cv2.matchTemplate） ===
TM_CCORR: max_err=2.240000e+02, mean_err=3.329635e+01, rel_max=8.109e-07, rel_mean=1.205e-07
TM_CCORR_NORMED: max_err=1.788139e-06, mean_err=1.768073e-07, rel_max=1.788e-06, rel_mean=1.768e-07
TM_SQDIFF: max_err=4.800000e+02, mean_err=6.757257e+01, rel_max=3.034e-06, rel_mean=4.272e-07
TM_SQDIFF_NORMED: max_err=1.142450e+00, mean_err=3.005487e-02, rel_max=1.142e+00, rel_mean=3.005e-02
TM_CCOEFF: max_err=2.491250e+02, mean_err=3.909091e+01, rel_max=2.171e-05, rel_mean=3.407e-06
TM_CCOEFF_NORMED: max_err=2.512969e-04, mean_err=4.180090e-06, rel_max=2.513e-04, rel_mean=4.180e-06

import numpy as np
import cv2
import matplotlib.pyplot as plt
import argparse
from pathlib import Path

def fft_xcorr2d_full(a, b):
    """
    full 互相关 a ⋆ b，尺寸 (Ha+Hb-1, Wa+Wb-1)
    实现：ifft2( FFT(a_pad) * conj(FFT(b_pad)) )
    """
    Ha, Wa = a.shape
    Hb, Wb = b.shape
    Hf, Wf = Ha + Hb - 1, Wa + Wb - 1
    Fa = np.fft.rfft2(a, s=(Hf, Wf))
    Fb = np.fft.rfft2(b, s=(Hf, Wf))
    Fc = Fa * np.conj(Fb)
    c_full = np.fft.irfft2(Fc, s=(Hf, Wf))
    return c_full

def valid_from_full_xcorr(full, H, W, h, w):
    """
    对于“互相关”的 full 结果，valid（窗口左上角对齐）应取：
    [0:H-h+1, 0:W-w+1]
    """
    return full[0:H - h + 1, 0:W - w + 1]

def compute_fft_stats(I, T):
    """
    用同一个坐标系（互相关 + 顶左对齐 valid）计算：
    - C = I ⋆ T
    - S1 = I 与 Ones 的相关（窗口和）
    - S2 = I^2 与 Ones 的相关（窗口平方和）
    返回 C, S1, S2 以及模板常量 ST, ST2, var_T, N
    """
    H, W = I.shape
    h, w = T.shape
    N = float(h * w)

    # 分子：互相关
    C_full = fft_xcorr2d_full(I, T)
    C = valid_from_full_xcorr(C_full, H, W, h, w)

    # 同样用互相关计算窗口统计量，确保对齐一致
    Ones = np.ones_like(T, dtype=I.dtype)
    S1_full = fft_xcorr2d_full(I, Ones)
    S1 = valid_from_full_xcorr(S1_full, H, W, h, w)

    S2_full = fft_xcorr2d_full(I * I, Ones)
    S2 = valid_from_full_xcorr(S2_full, H, W, h, w)

    ST = float(T.sum())
    ST2 = float((T * T).sum())
    var_T = ST2 - (ST * ST) / N

    return C, S1, S2, ST, ST2, var_T, N

def match_template_fft(image, template, use_float32=True):
    """
    六种方法（与 OpenCV 定义一致），全部基于“频域互相关 + 顶左对齐 valid”
    """
    dtype = np.float32 if use_float32 else np.float64
    I = image.astype(dtype, copy=False)
    T = template.astype(dtype, copy=False)

    C, S1, S2, ST, ST2, var_T, N = compute_fft_stats(I, T)

    # 1) CCORR
    R_CCORR = C

    # 2) CCORR_NORMED
    denom_ccorr = np.sqrt(S2 * ST2).astype(dtype)
    R_CCORR_NORMED = np.divide(
        R_CCORR, denom_ccorr, out=np.zeros_like(R_CCORR), where=denom_ccorr > 0
    )

    # 3) SQDIFF
    R_SQDIFF = S2 + ST2 - 2.0 * C

    # 4) SQDIFF_NORMED
    denom_sqdiff = np.sqrt(S2 * ST2).astype(dtype)
    R_SQDIFF_NORMED = np.divide(
        R_SQDIFF, denom_sqdiff, out=np.zeros_like(R_SQDIFF), where=denom_sqdiff > 0
    )

    # 5) CCOEFF
    R_CCOEFF = C - (S1 * ST) / dtype(N)

    # 6) CCOEFF_NORMED
    var_I = S2 - (S1 * S1) / dtype(N)
    var_I = np.maximum(var_I, 0.0).astype(dtype)
    var_T_safe = max(var_T, 0.0)
    denom_ccoef = np.sqrt(var_I * dtype(var_T_safe)).astype(dtype)
    R_CCOEFF_NORMED = np.divide(
        R_CCOEFF, denom_ccoef, out=np.zeros_like(R_CCOEFF), where=denom_ccoef > 0
    )

    return {
        "TM_CCORR": R_CCORR,
        "TM_CCORR_NORMED": R_CCORR_NORMED,
        "TM_SQDIFF": R_SQDIFF,
        "TM_SQDIFF_NORMED": R_SQDIFF_NORMED,
        "TM_CCOEFF": R_CCOEFF,
        "TM_CCOEFF_NORMED": R_CCOEFF_NORMED
    }

def compare_with_opencv(image, template):
    methods = {
        "TM_CCORR": cv2.TM_CCORR,
        "TM_CCORR_NORMED": cv2.TM_CCORR_NORMED,
        "TM_SQDIFF": cv2.TM_SQDIFF,
        "TM_SQDIFF_NORMED": cv2.TM_SQDIFF_NORMED,
        "TM_CCOEFF": cv2.TM_CCOEFF,
        "TM_CCOEFF_NORMED": cv2.TM_CCOEFF_NORMED,
    }

    ours = match_template_fft(image, template, use_float32=True)

    report = {}
    diffs = {}
    opencv_maps = {}
    for name, m in methods.items():
        opencv_res = cv2.matchTemplate(image.astype(np.float32),
                                       template.astype(np.float32),
                                       m)
        ours_map = ours[name].astype(opencv_res.dtype, copy=False)
        diff = ours_map - opencv_res
        max_err = float(np.max(np.abs(diff)))
        mean_err = float(np.mean(np.abs(diff)))
        scale = float(max(1.0, np.max(np.abs(opencv_res))))
        rel_max = max_err / scale
        rel_mean = mean_err / scale
        report[name] = {
            "opencv_shape": opencv_res.shape,
            "our_shape": ours_map.shape,
            "max_abs_error": max_err,
            "mean_abs_error": mean_err,
            "rel_max_error": rel_max,
            "rel_mean_error": rel_mean
        }
        diffs[name] = diff
        opencv_maps[name] = opencv_res
    return report, ours, opencv_maps, diffs

def draw_results(image, template, ours, opencv_maps, diffs, src_top_left, save_path=None):
    """
    可视化：
    - 每种方法一行：ours / opencv / |diff|
    - 另绘原图+最佳匹配框（用 CCOEFF_NORMED 和 SQDIFF_NORMED）
    """
    methods_order = [
        "TM_CCORR",
        "TM_CCORR_NORMED",
        "TM_SQDIFF",
        "TM_SQDIFF_NORMED",
        "TM_CCOEFF",
        "TM_CCOEFF_NORMED",
    ]

    nrows, ncols = len(methods_order), 3
    fig, axes = plt.subplots(nrows, ncols, figsize=(12, 2.8*nrows), constrained_layout=True)
    if nrows == 1:
        axes = np.array([axes])

    for i, name in enumerate(methods_order):
        ax1, ax2, ax3 = axes[i]

        m1 = ours[name]
        m2 = opencv_maps[name]
        d = np.abs(diffs[name])

        im1 = ax1.imshow(m1, cmap="viridis")
        ax1.set_title(f"{name} (ours)")
        fig.colorbar(im1, ax=ax1, fraction=0.046, pad=0.04)

        im2 = ax2.imshow(m2, cmap="viridis")
        ax2.set_title(f"{name} (opencv)")
        fig.colorbar(im2, ax=ax2, fraction=0.046, pad=0.04)

        im3 = ax3.imshow(d, cmap="magma")
        ax3.set_title(f"|diff| (ours - opencv)")
        fig.colorbar(im3, ax=ax3, fraction=0.046, pad=0.04)

        for ax in (ax1, ax2, ax3):
            ax.set_xticks([])
            ax.set_yticks([])

    if save_path:
        fig.savefig(save_path[0], dpi=150, bbox_inches='tight')
        print(f"Saved: {save_path[0]}")
    else:
        plt.show()
    plt.close(fig)

    # 在原图上画最佳匹配（以 CCOEFF_NORMED 最大值和 SQDIFF_NORMED 最小值为例）
    H, W = image.shape
    h, w = template.shape

    def draw_rect_on_image(ax, top_left, color='lime', label=''):
        r0, c0 = top_left
        ax.add_patch(plt.Rectangle((c0, r0), w, h, fill=False, edgecolor=color, linewidth=2))
        if label:
            ax.text(c0, r0-5, label, color=color, fontsize=10, backgroundcolor='k')

    best_maps = [
        ("TM_CCOEFF_NORMED", True, "max"),  # 最大越好
        ("TM_SQDIFF_NORMED", False, "min"), # 最小越好
    ]

    fig2, ax2 = plt.subplots(1, 2, figsize=(12, 5), constrained_layout=True)
    for k, (name, is_max, tag) in enumerate(best_maps):
        if is_max:
            best_pos_ours = np.unravel_index(int(np.argmax(ours[name])), ours[name].shape)
            best_pos_cv = np.unravel_index(int(np.argmax(opencv_maps[name])), opencv_maps[name].shape)
        else:
            best_pos_ours = np.unravel_index(int(np.argmin(ours[name])), ours[name].shape)
            best_pos_cv = np.unravel_index(int(np.argmin(opencv_maps[name])), opencv_maps[name].shape)

        ax = ax2[k]
        ax.imshow(image, cmap='gray')
        draw_rect_on_image(ax, (int(best_pos_ours[0]), int(best_pos_ours[1])), color='lime', label=f"ours {tag}")
        draw_rect_on_image(ax, (int(best_pos_cv[0]), int(best_pos_cv[1])), color='cyan', label=f"opencv {tag}")
        draw_rect_on_image(ax, (int(src_top_left[0]), int(src_top_left[1])), color='yellow', label="template src")
        ax.set_title(f"Best by {name}")
        ax.set_xticks([]); ax.set_yticks([])

    if save_path and len(save_path) > 1:
        fig2.savefig(save_path[1], dpi=150, bbox_inches='tight')
        print(f"Saved: {save_path[1]}")
    else:
        plt.show()
    plt.close(fig2)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", type=str, default="test.png", help="输入灰度图路径，默认 test.png")
    parser.add_argument("--x", type=int, default=None, help="模板左上角 x（行）")
    parser.add_argument("--y", type=int, default=None, help="模板左上角 y（列）")
    parser.add_argument("--h", type=int, default=None, help="模板高度")
    parser.add_argument("--w", type=int, default=None, help="模板宽度")
    args = parser.parse_args()

    img_path = Path(args.image)
    if not img_path.exists():
        raise FileNotFoundError(f"找不到图像文件: {img_path}")

    # 读灰度
    image = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE)
    if image is None:
        raise RuntimeError("读取图像失败，请确认路径与文件格式。")
    image = image.astype(np.float32)

    H, W = image.shape

    # 默认模板：图像中部区域
    h = args.h if args.h is not None else max(16, H // 6)
    w = args.w if args.w is not None else max(16, W // 6)
    x0 = args.x if args.x is not None else max(0, (H // 3))
    y0 = args.y if args.y is not None else max(0, (W // 3))

    # 防越界
    x0 = min(max(0, x0), max(0, H - h))
    y0 = min(max(0, y0), max(0, W - w))

    template = image[x0:x0+h, y0:y0+w].copy()

    # 计算与对比
    report, ours, opencv_maps, diffs = compare_with_opencv(image, template)

    print("=== 数值对比（与 cv2.matchTemplate） ===")
    for k, v in report.items():
        print(f"{k}: max_err={v['max_abs_error']:.6e}, mean_err={v['mean_abs_error']:.6e}, "
              f"rel_max={v['rel_max_error']:.3e}, rel_mean={v['rel_mean_error']:.3e}")

    # 归一化方法应在合理范围（CCOEFF_NORMED、CCORR_NORMED ∈ [-1,1]；SQDIFF_NORMED ∈ [0,2]）
    for name in ["TM_CCORR_NORMED", "TM_SQDIFF_NORMED", "TM_CCOEFF_NORMED"]:
        m = ours[name]
        print(f"{name} range: min={float(m.min()):.6f}, max={float(m.max()):.6f}")

    # 可视化
    draw_results(image, template, ours, opencv_maps, diffs, src_top_left=(x0, y0),
                 save_path=["matching_results.png", "best_match_locations.png"])

if __name__ == "__main__":
    main()