hello_js_reverse_skill/tcaptcha/solvers/icon_click.py

# -*- coding: utf-8 -*-
"""
腾讯天御点选验证码工业级纯 ONNX 视觉求解器 (Solver v3.0 - Pure ONNX Edition)
特性:
1. 零 PyTorch / 零 Ultralytics 依赖：采用 onnxruntime 推理引擎 (从 1.5GB 骤降至 15MB)
2. 提示词解析: 垂直/水平投影切片 + ResNet18 ONNX 分类器直接预测目标序列
3. 候选图标定位: YOLOv8n ONNX + Letterbox 自适应缩放 + OpenCV cv2.dnn.NMSBoxes 精准去重
4. 全局匹配: 后验概率分布优化 + 背景干扰项抑制因子 + 全局非冲突组合搜索，输出有序像素点击坐标
"""

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

import cv2
import numpy as np
import onnxruntime as ort
from PIL import Image

from tcaptcha.ocr.slicer import slice_prompt_image

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

# 权重在包根的 models/（本文件在 tcaptcha/solvers/ 下，故上跳一层）
MODEL_DIR = Path(__file__).resolve().parents[1] / "models"
YOLO_ONNX_PATH = MODEL_DIR / "yolo_detector.onnx"
CLASSIFIER_ONNX_PATH = MODEL_DIR / "icon_classifier.onnx"


class CaptchaSolver:
    def __init__(self, yolo_path: Path = YOLO_ONNX_PATH, classifier_path: Path = CLASSIFIER_ONNX_PATH):
        if not yolo_path.exists():
            raise FileNotFoundError(f"YOLO ONNX 权重文件不存在: {yolo_path}")
        if not classifier_path.exists():
            raise FileNotFoundError(f"Classifier ONNX 权重文件不存在: {classifier_path}")

        # 配置轻量级 CPU 推理会话
        sess_options = ort.SessionOptions()
        sess_options.inter_op_num_threads = 2
        sess_options.intra_op_num_threads = 2
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

        self.yolo_sess = ort.InferenceSession(str(yolo_path), sess_options, providers=["CPUExecutionProvider"])
        self.cls_sess = ort.InferenceSession(str(classifier_path), sess_options, providers=["CPUExecutionProvider"])

    @staticmethod
    def _preprocess_patch(img: Image.Image) -> np.ndarray:
        """对小图标图像进行 64x64 归一化预处理 (匹配 ResNet18 输入标准)"""
        resized = img.convert("RGB").resize((64, 64), Image.Resampling.BILINEAR)
        arr = np.array(resized, dtype=np.float32) / 255.0
        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
        norm = (arr - mean) / std
        return np.transpose(norm, (2, 0, 1))  # (3, 64, 64)

    def parse_prompt_targets(self, prompt_img: Union[str, Path, bytes, Image.Image]) -> List[int]:
        """自动切片并解析提示图中 3 个目标图标的类别 ID (0~11)"""
        crops = slice_prompt_image(prompt_img)
        if not crops:
            return []

        patch_blobs = np.stack([self._preprocess_patch(c[0]) for c in crops[:3]], axis=0).astype(np.float32)
        logits = self.cls_sess.run(None, {"input": patch_blobs})[0]  # shape: (N, 13)
        return [int(np.argmax(l[:12])) for l in logits]

    def _extract_candidates(self, bg_img: Union[str, Path, bytes, Image.Image]) -> List[Tuple[int, int, Image.Image]]:
        """利用 YOLOv8n ONNX 召回背景图所有图标候选位置并去重"""
        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")

        orig_w, orig_h = bg_pil.size
        img_np = np.array(bg_pil)

        # 1. Letterbox 保持宽高比填充缩放至 640x640
        r = min(640.0 / orig_h, 640.0 / orig_w)
        new_w, new_h = int(round(orig_w * r)), int(round(orig_h * r))
        resized = cv2.resize(img_np, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        dw = (640.0 - new_w) / 2.0
        dh = (640.0 - new_h) / 2.0
        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
        padded = cv2.copyMakeBorder(resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))

        blob = padded.astype(np.float32) / 255.0
        blob = np.transpose(blob, (2, 0, 1))[None, ...]

        # 2. ONNX 目标检测前向推理
        yolo_out = self.yolo_sess.run(None, {"images": blob})[0][0].T  # (8400, 5)
        valid = yolo_out[yolo_out[:, 4] >= 0.03]

        boxes_xywh = []
        scores = []
        for row in valid:
            cx, cy, w, h, conf = row
            x1 = (cx - w / 2.0 - dw) / r
            y1 = (cy - h / 2.0 - dh) / r
            boxes_xywh.append([int(x1), int(y1), int(w / r), int(h / r)])
            scores.append(float(conf))

        # 3. OpenCV NMS 筛选重叠边界框
        indices = cv2.dnn.NMSBoxes(boxes_xywh, scores, score_threshold=0.03, nms_threshold=0.45)
        raw_cands = []
        if len(indices) > 0:
            for idx in indices:
                x, y, w, h = boxes_xywh[idx]
                raw_cands.append((int(x + w / 2.0), int(y + h / 2.0)))

        # 4. 空间中心点近邻聚类去重 (<=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(orig_w, cx + 32), min(orig_h, cy + 32))
            patch = bg_pil.crop(crop_box)
            candidates.append((cx, cy, patch))

        return candidates

    def solve(
        self,
        bg_img: Union[str, Path, bytes, Image.Image],
        prompt_img: Union[str, Path, bytes, Image.Image],
    ) -> List[Tuple[int, int]]:
        """
        端到端纯 ONNX 求解验证码点选坐标
        返回 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. 批量推理计算候选图标的类别概率分布 (13 维)
        patch_blobs = np.stack([self._preprocess_patch(c[2]) for c in candidates], axis=0).astype(np.float32)
        logits = self.cls_sess.run(None, {"input": patch_blobs})[0]
        exp = np.exp(logits - np.max(logits, axis=1, keepdims=True))
        probs = exp / np.sum(exp, axis=1, keepdims=True)

        # 4. 全局后验概率优化 (结合背景干扰项抑制因子 1 - P(distractor))
        bg_probs = probs[:, 12]
        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()
            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))
                comb_list.append(best_c)
                used.add(best_c)
            best_comb = tuple(comb_list)

        # 5. 输出有序点击坐标
        return [(candidates[idx][0], candidates[idx][1]) for idx in best_comb]


if __name__ == "__main__":
    solver = CaptchaSolver()
    print("✓ CaptchaSolver ONNX 引擎初始化成功！")