douyin-captcha/try/method_d_shape.py

"""方法D：基于轮廓/边缘的缩放旋转搜索。

mark 的 RGB 颜色可能与大图中的目标渲染颜色不同，因此不再匹配填充颜色，
而只使用 mark alpha 的外轮廓；背景使用 Canny 边缘图，采用距离变换做
Chamfer matching。分数越高表示模板轮廓落在背景边缘上的平均距离越小。

输出：try/out/D/{*-match.png,summary.txt}
"""
from pathlib import Path
import math
import cv2
import numpy as np

ROOT = Path(__file__).resolve().parent.parent
CAP = ROOT / "captchas"
OUT = ROOT / "try" / "out" / "D"
OUT.mkdir(parents=True, exist_ok=True)

SCALES = np.arange(0.30, 1.51, 0.05)
ANGLES = np.arange(-180.0, 180.0, 15.0)
TOP_K = 5


def files_for(prefix):
    bigs = [p for p in CAP.glob(f"{prefix}*.jpeg") if "-mark" not in p.stem]
    if not bigs:
        return None, None
    big = bigs[0]
    return big, CAP / f"{big.stem}-mark.png"


def load_shapes(mark_path, big_path):
    try:
        mark = cv2.imread(str(mark_path), cv2.IMREAD_UNCHANGED)
        big = cv2.imread(str(big_path), cv2.IMREAD_GRAYSCALE)
        if mark is None or big is None:
            return None
        alpha = mark[..., 3] if mark.ndim == 3 and mark.shape[2] == 4 else np.full(mark.shape[:2], 255, np.uint8)
        shape = (alpha > 128).astype(np.uint8) * 255
        # 轮廓用细边界，避免大面积填充主导匹配
        edge = cv2.morphologyEx(shape, cv2.MORPH_GRADIENT, np.ones((3, 3), np.uint8))
        return big, shape, edge
    except Exception as e:
        print(f"  图像读取失败: {e}")
        return None


def transformed_edge(edge, scale, angle):
    try:
        h0, w0 = edge.shape
        w, h = max(12, round(w0 * scale)), max(12, round(h0 * scale))
        small = cv2.resize(edge, (w, h), interpolation=cv2.INTER_NEAREST)
        if abs(angle) < 1e-6:
            return small
        m = cv2.getRotationMatrix2D((w / 2.0, h / 2.0), angle, 1.0)
        cos, sin = abs(float(m[0, 0])), abs(float(m[0, 1]))
        nw, nh = int(h * sin + w * cos) + 1, int(h * cos + w * sin) + 1
        m[0, 2] += nw / 2.0 - w / 2.0
        m[1, 2] += nh / 2.0 - h / 2.0
        return cv2.warpAffine(small, m, (nw, nh), flags=cv2.INTER_NEAREST, borderValue=0)
    except cv2.error:
        return None


def chamfer_response(bg_edges, tpl_edge):
    """边缘模板的 Chamfer 响应图，值越大越好。"""
    try:
        h, w = tpl_edge.shape
        if h >= bg_edges.shape[0] or w >= bg_edges.shape[1]:
            return None
        points = (tpl_edge > 0).astype(np.float32)
        n_points = float(points.sum())
        if n_points < 10:
            return None
        # 用模板边缘作为权重，计算背景边缘距离的局部平均值。
        distance = cv2.distanceTransform((255 - bg_edges).astype(np.uint8), cv2.DIST_L2, 3)
        cost = cv2.matchTemplate(distance.astype(np.float32), points,
                                 cv2.TM_CCORR)
        avg_distance = cost / n_points
        return 1.0 / (1.0 + np.maximum(avg_distance, 0.0))
    except cv2.error:
        return None


def peaks(response, tw, th):
    out = []
    try:
        r = response.copy()
        for _ in range(TOP_K):
            _, score, _, loc = cv2.minMaxLoc(r)
            if score <= 0:
                break
            out.append((float(score), int(loc[0]), int(loc[1])))
            x0 = max(0, loc[0] - int(tw * 0.65))
            y0 = max(0, loc[1] - int(th * 0.65))
            x1 = min(r.shape[1], loc[0] + int(tw * 0.65))
            y1 = min(r.shape[0], loc[1] + int(th * 0.65))
            r[y0:y1, x0:x1] = 0
    except Exception as e:
        print(f"  峰提取失败: {e}")
    return out


def process(prefix, lines):
    try:
        big_path, mark_path = files_for(prefix)
        if big_path is None:
            lines.append(f"{prefix}\tSKIP\t文件缺失")
            return
        loaded = load_shapes(mark_path, big_path)
        if loaded is None:
            lines.append(f"{prefix}\tSKIP\t读取失败")
            return
        big, _, edge = loaded
        bg_edges = cv2.Canny(big, 50, 150)
        candidates = []
        for scale in SCALES:
            for angle in ANGLES:
                tpl = transformed_edge(edge, float(scale), float(angle))
                if tpl is None:
                    continue
                response = chamfer_response(bg_edges, tpl)
                if response is None:
                    continue
                for score, x, y in peaks(response, tpl.shape[1], tpl.shape[0]):
                    candidates.append((score, float(scale), float(angle), x, y,
                                      tpl.shape[1], tpl.shape[0]))
        if not candidates:
            lines.append(f"{prefix}\tSKIP\t无候选")
            return
        # 跨尺度/角度 NMS
        candidates.sort(reverse=True)
        kept = []
        for item in candidates:
            score, scale, angle, x, y, w, h = item
            cx, cy = x + w / 2, y + h / 2
            if all(math.hypot(cx - (q[3] + q[5] / 2), cy - (q[4] + q[6] / 2)) > 0.6 * max(w, h)
                   for q in kept):
                kept.append(item)
            if len(kept) >= TOP_K:
                break
        best = kept[0]
        score, scale, angle, x, y, w, h = best
        print(f"{prefix}  best score={score:.3f} scale={scale:.2f} angle={angle:+.0f} "
              f"offset=({x},{y}) top={len(kept)}")
        lines.append(f"{prefix}\tscore={score:.3f}\tscale={scale:.2f}\tangle={angle:+.0f}\t"
                     f"offset=({x},{y})\ttop={len(kept)}")
        vis = cv2.cvtColor(big, cv2.COLOR_GRAY2BGR)
        colors = [(0, 0, 255), (0, 165, 255), (255, 0, 0), (0, 255, 0), (255, 0, 255)]
        for i, item in enumerate(kept):
            _, _, _, px, py, pw, ph = item
            cv2.rectangle(vis, (px, py), (px + pw, py + ph), colors[i], 2)
            cv2.putText(vis, str(i + 1), (px + 2, py + 16), cv2.FONT_HERSHEY_SIMPLEX,
                        0.55, colors[i], 2)
        cv2.imwrite(str(OUT / f"{prefix}-match.png"), vis)
    except Exception as e:
        print(f"{prefix}  处理失败: {e}")
        lines.append(f"{prefix}\tERROR\t{e}")


def main():
    lines = []
    prefixes = sorted({p.name.split("~", 1)[0] for p in CAP.glob("*.jpeg")
                       if "-mark" not in p.stem})
    for prefix in prefixes:
        process(prefix, lines)
    (OUT / "summary.txt").write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"\n结果已写入 {OUT}/")


if __name__ == "__main__":
    main()