douyin-captcha/try/method_e_shape_coverage.py

"""方法E：轮廓边缘覆盖率匹配。

与方法D不同，不直接比较边缘距离的原始平均值（小模板容易偶然得高分），
而统计模板轮廓点落在背景边缘容差范围内的比例，并以平均距离作为次级指标。
候选按 (coverage, -mean_distance) 排序。

输出：try/out/E/{*-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" / "E"
OUT.mkdir(parents=True, exist_ok=True)

SCALES = np.arange(0.40, 1.31, 0.05)
ANGLES = np.arange(-180.0, 180.0, 15.0)
TOP_K = 3
EDGE_TOL = 2.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_shape(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)
        mask = (alpha > 128).astype(np.uint8) * 255
        edge = cv2.morphologyEx(mask, cv2.MORPH_GRADIENT, np.ones((3, 3), np.uint8))
        return big, edge
    except Exception as e:
        print(f"  读取失败: {e}")
        return None


def transform_edge(edge, scale, angle):
    try:
        h0, w0 = edge.shape
        w, h = max(16, round(w0 * scale)), max(16, 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)
        c, s = abs(float(m[0, 0])), abs(float(m[0, 1]))
        nw, nh = int(h * s + w * c) + 1, int(h * c + w * s) + 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 response_maps(bg_edges):
    try:
        dist = cv2.distanceTransform((255 - bg_edges).astype(np.uint8), cv2.DIST_L2, 3)
        near = (dist <= EDGE_TOL).astype(np.float32)
        return dist.astype(np.float32), near
    except cv2.error:
        return None, None


def candidates_for_template(dist, near, tpl):
    try:
        points = (tpl > 0).astype(np.float32)
        n = float(points.sum())
        h, w = tpl.shape
        if n < 20 or h >= dist.shape[0] or w >= dist.shape[1]:
            return []
        hit = cv2.matchTemplate(near, points, cv2.TM_CCORR) / n
        cost = cv2.matchTemplate(dist, points, cv2.TM_CCORR) / n
        out = []
        for _ in range(TOP_K):
            _, cov, _, loc = cv2.minMaxLoc(hit)
            mean_dist = float(cost[loc[1], loc[0]])
            if cov <= 0:
                break
            out.append((float(cov), -mean_dist, int(loc[0]), int(loc[1])))
            x0 = max(0, loc[0] - int(w * 0.7))
            y0 = max(0, loc[1] - int(h * 0.7))
            x1 = min(hit.shape[1], loc[0] + int(w * 0.7))
            y1 = min(hit.shape[0], loc[1] + int(h * 0.7))
            hit[y0:y1, x0:x1] = 0
        return out
    except cv2.error:
        return []


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_shape(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)
        dist, near = response_maps(bg_edges)
        if dist is None or near is None:
            lines.append(f"{prefix}\tSKIP\t无法生成距离图")
            return
        all_candidates = []
        for scale in SCALES:
            for angle in ANGLES:
                tpl = transform_edge(edge, float(scale), float(angle))
                if tpl is None:
                    continue
                for cov, neg_dist, x, y in candidates_for_template(dist, near, tpl):
                    all_candidates.append((cov, neg_dist, float(scale), float(angle), x, y,
                                           tpl.shape[1], tpl.shape[0]))
        if not all_candidates:
            lines.append(f"{prefix}\tSKIP\t无候选")
            return
        all_candidates.sort(key=lambda x: (-x[0], -x[1]))
        kept = []
        for item in all_candidates:
            cov, neg_dist, scale, angle, x, y, w, h = item
            cx, cy = x + w / 2.0, y + h / 2.0
            if all(math.hypot(cx - (q[4] + q[6] / 2), cy - (q[5] + q[7] / 2)) > 0.6 * max(w, h)
                   for q in kept):
                kept.append(item)
            if len(kept) >= TOP_K:
                break
        best = kept[0]
        cov, neg_dist, scale, angle, x, y, w, h = best
        print(f"{prefix}  coverage={cov:.3f} meanDist={-neg_dist:.2f} "
              f"scale={scale:.2f} angle={angle:+.0f} offset=({x},{y})")
        lines.append(f"{prefix}\tcoverage={cov:.3f}\tmeanDist={-neg_dist:.2f}\t"
                     f"scale={scale:.2f}\tangle={angle:+.0f}\toffset=({x},{y})")
        vis = cv2.cvtColor(big, cv2.COLOR_GRAY2BGR)
        colors = [(0, 0, 255), (0, 165, 255), (255, 0, 0)]
        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()