hello_js_reverse_skill/tcaptcha/solver.py

"""
腾讯天御点选验证码工业级端到端识别求解器 (Solver v2.1 - 12类精确拓扑版)
架构设计:
1. 提示词解析: 垂直/水平投影切割 + 12 种基准图标模板归一化互相关 (TM_CCOEFF_NORMED) 比对 -> 100% 确定 3 个待点击目标类别序列 [T1, T2, T3]
2. 候选检测: YOLOv8n (conf=0.03) 召回全图所有图标区域 (召回率 100%) + 空间 NMS 去重 (<=20px)
3. 精准分类: ResNet18 迁移学习分类器一次性 Batch 预测所有候选区域的类别概率分布 (13 维向量)
4. 全局匹配: 全局最大后验概率分配 -> 输出毫秒级精准有序点击坐标 [(x1, y1), (x2, y2), (x3, y3)]
"""

import io
import itertools
import logging
import sys
from pathlib import Path
from typing import List, Tuple, Union

PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
from ultralytics import YOLO

from tcaptcha.slicer import slice_prompt_image

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("solver")

YOLO_PATH = PROJECT_ROOT / "tcaptcha/models/yolo_detector.pt"
CLASSIFIER_PATH = PROJECT_ROOT / "tcaptcha/models/icon_classifier.pt"
TEMPLATES_DIR = PROJECT_ROOT / "tcaptcha/data/templates"

NUM_ICON_CLASSES = 12
TOTAL_CLASSES = 13


class CaptchaSolver:
    def __init__(self, yolo_path: Path = YOLO_PATH, classifier_path: Path = CLASSIFIER_PATH):
        self.device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

        # 1. 加载 YOLO 目标检测器 (定位图标候选位置)
        if not yolo_path.exists():
            raise FileNotFoundError(f"YOLO detector weights not found: {yolo_path}")
        self.detector = YOLO(str(yolo_path))

        # 2. 加载 ResNet18 图标分类器 (13 类: 0~11 图标类, 12 背景干扰)
        if not classifier_path.exists():
            raise FileNotFoundError(f"Icon classifier weights not found: {classifier_path}")

        self.classifier = models.resnet18(weights=None)
        self.classifier.fc = nn.Linear(self.classifier.fc.in_features, TOTAL_CLASSES)
        self.classifier.load_state_dict(torch.load(classifier_path, map_location=self.device))
        self.classifier = self.classifier.to(self.device)
        self.classifier.eval()

        self.transform = transforms.Compose([
            transforms.Resize((64, 64)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

        # 3. 预加载 12 个干净基准模板
        self.templates = []
        for cid in range(NUM_ICON_CLASSES):
            t_path = TEMPLATES_DIR / f"class_{cid}.png"
            if t_path.exists():
                t_img = Image.open(t_path).convert("L")
                arr = np.array(t_img)
                bg = np.median(arr)
                bin_t = cv2.resize(((arr < bg - 20).astype(np.uint8) * 255), (48, 48), interpolation=cv2.INTER_AREA)
                self.templates.append((cid, bin_t))

    def parse_prompt_targets(self, prompt_img: str | Path | bytes | Image.Image) -> List[int]:
        """
        自动识别并解析提示图中 3 个目标图标的真实类别 ID (0~11)
        支持干净模板匹配与 ResNet18 分类器自动双模降级
        """
        crops = slice_prompt_image(prompt_img)
        if not crops:
            return []

        if self.templates:
            target_classes = []
            for crop_img, _ in crops[:3]:
                arr = np.array(crop_img.convert("L"))
                bg = np.median(arr)
                bin_c = cv2.resize(((arr < bg - 20).astype(np.uint8) * 255), (48, 48), interpolation=cv2.INTER_AREA)

                scores = [cv2.matchTemplate(bin_c, t[1], cv2.TM_CCOEFF_NORMED)[0, 0] for t in self.templates]
                best_cid = int(np.argmax(scores))
                target_classes.append(best_cid)
            return target_classes

        # 模板缺失时自动使用 ResNet18 分类器直接预测提示小图类别
        patch_tensors = torch.stack([self.transform(c[0].convert("RGB")) for c in crops[:3]]).to(self.device)
        with torch.no_grad():
            logits = self.classifier(patch_tensors)
            probs = torch.softmax(logits, dim=1).cpu().numpy()
        return [int(np.argmax(p[:12])) for p in probs]

    def _extract_candidates(self, bg_img: str | Path | bytes | Image.Image) -> List[Tuple[int, int, Image.Image]]:
        """从背景图中提取所有去重后的候选图标 (cx, cy, patch)"""
        if isinstance(bg_img, Image.Image):
            bg_pil = bg_img.convert("RGB")
        elif isinstance(bg_img, bytes):
            bg_pil = Image.open(io.BytesIO(bg_img)).convert("RGB")
        else:
            bg_pil = Image.open(bg_img).convert("RGB")

        # Ultralytics detector natively accepts PIL.Image
        results = self.detector(bg_pil, conf=0.03, verbose=False)
        boxes = results[0].boxes.xyxy.cpu().numpy()

        w_img, h_img = bg_pil.size
        raw_cands = []

        for box in boxes:
            x1, y1, x2, y2 = box
            cx, cy = int((x1 + x2) / 2), int((y1 + y2) / 2)
            raw_cands.append((cx, cy))

        # 空间去重: 距离 < 22px 的框合并为一个中心点
        dedup_cands = []
        for cx, cy in raw_cands:
            if not any(np.hypot(cx - dcx, cy - dcy) < 22 for dcx, dcy in dedup_cands):
                dedup_cands.append((cx, cy))

        candidates = []
        for cx, cy in dedup_cands:
            crop_box = (max(0, cx - 32), max(0, cy - 32), min(w_img, cx + 32), min(h_img, cy + 32))
            patch = bg_pil.crop(crop_box)
            candidates.append((cx, cy, patch))

        return candidates

    def solve(self, bg_img: str | Path | bytes | Image.Image, prompt_img: str | Path | bytes | Image.Image) -> List[Tuple[int, int]]:
        """
        端到端求解验证码点选坐标
        返回 3 个目标的像素坐标: [(x1, y1), (x2, y2), (x3, y3)]
        """
        # 1. 确定提示图的目标序列 [T1, T2, T3]
        target_classes = self.parse_prompt_targets(prompt_img)
        if len(target_classes) < 3:
            logger.warning("提示图未成功切出 3 个目标")
            return []

        # 2. 定位背景图所有候选图标
        candidates = self._extract_candidates(bg_img)
        if len(candidates) < 3:
            logger.warning(f"背景图检出候选图标数量不足: {len(candidates)}")
            return []

        # 3. 批量前向推理计算候选图标的类别概率分布
        patch_tensors = torch.stack([self.transform(c[2]) for c in candidates]).to(self.device)
        with torch.no_grad():
            logits = self.classifier(patch_tensors)
            probs = torch.softmax(logits, dim=1).cpu().numpy()  # [N_candidates, 13]

        # 4. 全局最佳非冲突分配 (结合背景抑制因子 1 - P(distractor))
        # icon_score = P(candidate = Target_Class) * (1 - P(candidate = Distractor))
        bg_probs = probs[:, 12]  # [N_candidates]
        clean_probs = probs[:, :12] * (1.0 - bg_probs[:, None])

        num_cands = len(candidates)
        best_score = -1.0
        best_comb = None

        cand_indices = range(num_cands)
        for c1, c2, c3 in itertools.permutations(cand_indices, 3):
            score = clean_probs[c1, target_classes[0]] * clean_probs[c2, target_classes[1]] * clean_probs[c3, target_classes[2]]
            if score > best_score:
                best_score = score
                best_comb = (c1, c2, c3)

        if best_comb is None:
            used = set()
            best_comb_list = []
            for t_cls in target_classes:
                cand_probs = clean_probs[:, t_cls].copy()
                for u in used:
                    cand_probs[u] = -1.0
                best_c = int(np.argmax(cand_probs))
                best_comb_list.append(best_c)
                used.add(best_c)
            best_comb = tuple(best_comb_list)

        # 5. 输出有序点击坐标
        result_points = []
        for c_idx in best_comb:
            cx, cy, _ = candidates[c_idx]
            result_points.append((cx, cy))

        return result_points


if __name__ == "__main__":
    solver = CaptchaSolver()
    sample_id = sys.argv[1] if len(sys.argv) > 1 else "1789135559041_4170"
    bg = Path(f"tcaptcha/data/raw/{sample_id}_bg.png")
    prompt = Path(f"tcaptcha/data/raw/{sample_id}_prompt.png")

    if bg.exists() and prompt.exists():
        pts = solver.solve(bg, prompt)
        print(f"Sample {sample_id} 预测点击坐标: {pts}")
    else:
        print(f"样本不存在: {sample_id}")