douyin-captcha/try/method_k_objectness.py

"""方法K：alpha 形状与暗色覆盖物概率图匹配。

不直接使用 mark RGB。根据样本结构，背景目标通常比周围照片更暗、饱和度
更低、局部纹理更均匀。构造暗色覆盖物概率图后，用 alpha silhouette 做
多尺度匹配，并通过局部峰值抑制避免小纹理重复命中。

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

SCALES = np.arange(0.65, 1.36, 0.025)
DARK_SIGMAS = (15.0, 21.0, 31.0)
SAT_WEIGHTS = (0.25, 0.5, 0.75)
TOP_K = 5


def pair_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 read_pair(big_path, mark_path):
    try:
        big = cv2.imread(str(big_path), cv2.IMREAD_COLOR)
        mark = cv2.imread(str(mark_path), cv2.IMREAD_UNCHANGED)
        if big is None or mark 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.float32)
        return big, shape
    except Exception as e:
        print(f"  读取失败: {e}")
        return None


def objectness_maps(bgr):
    """生成暗色覆盖物概率图的多个版本。"""
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY).astype(np.float32)
    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
    sat = hsv[..., 1] / 255.0
    maps = []
    for sigma in DARK_SIGMAS:
        try:
            local = cv2.GaussianBlur(gray, (0, 0), sigma)
            dark = np.maximum(local - gray, 0.0)
            # 局部变暗 + 低饱和度，抑制花田/山地的彩色纹理
            lo = float(dark.min())
            hi = float(dark.max())
            if hi > lo:
                dark = (dark - lo) / (hi - lo)
            else:
                dark = np.zeros_like(dark, dtype=np.float32)
            for sw in SAT_WEIGHTS:
                obj = dark * (1.0 - sw * sat)
                maps.append(obj.astype(np.float32))
        except cv2.error:
            continue
    return maps


def scaled_shape(shape, scale):
    try:
        h, w = shape.shape
        nw, nh = max(16, round(w * scale)), max(16, round(h * scale))
        return cv2.resize(shape, (nw, nh), interpolation=cv2.INTER_AREA).astype(np.float32)
    except cv2.error:
        return None


def response(objectness, shape):
    try:
        h, w = shape.shape
        if h >= objectness.shape[0] or w >= objectness.shape[1]:
            return None
        # 只在 alpha 前景内统计暗色覆盖程度
        return cv2.matchTemplate(objectness, shape, cv2.TM_CCORR_NORMED, mask=shape)
    except cv2.error:
        return None


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


def process(prefix, lines):
    try:
        big_path, mark_path = pair_for(prefix)
        if big_path is None:
            return
        loaded = read_pair(big_path, mark_path)
        if loaded is None:
            lines.append(f"{prefix}\tSKIP\t读取失败")
            return
        big, shape = loaded
        maps = objectness_maps(big)
        candidates = []
        for obj in maps:
            for scale in SCALES:
                tpl = scaled_shape(shape, float(scale))
                if tpl is None:
                    continue
                res = response(obj, tpl)
                if res is None:
                    continue
                for score, x, y in peaks(res, tpl.shape[1], tpl.shape[0]):
                    candidates.append((score, float(scale), x, y, tpl.shape[1], tpl.shape[0]))
        if not candidates:
            lines.append(f"{prefix}\tSKIP\t无候选")
            return
        candidates.sort(key=lambda c: -c[0])
        kept = []
        for c in candidates:
            score, scale, x, y, w, h = c
            cx, cy = x + w / 2.0, y + h / 2.0
            if all(math.hypot(cx - (q[2] + q[4] / 2), cy - (q[3] + q[5] / 2)) > 0.65 * max(w, h)
                   for q in kept):
                kept.append(c)
            if len(kept) >= TOP_K:
                break
        score, scale, x, y, w, h = kept[0]
        print(f"{prefix}  score={score:.3f} scale={scale:.3f} offset=({x},{y}) size=({w},{h}) top={len(kept)}")
        lines.append(f"{prefix}\tscore={score:.3f}\tscale={scale:.3f}\toffset=({x},{y})\t"
                     f"size=({w},{h})\ttop={len(kept)}")
        vis = big.copy()
        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, .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()