hello_js_reverse_skill/dist/rainyun_qinglong/solver.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
import sys
from pathlib import Path
from typing import List, Tuple, Union

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

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

try:
    from slicer import slice_prompt_image
except ImportError:
    from tcaptcha.slicer import slice_prompt_image

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

# 权重路径解析 (优先当前目录，兼容 models/ 子目录)
YOLO_ONNX_PATH = CURRENT_DIR / "yolo_detector.onnx"
if not YOLO_ONNX_PATH.exists():
    YOLO_ONNX_PATH = CURRENT_DIR / "models/yolo_detector.onnx"

CLASSIFIER_ONNX_PATH = CURRENT_DIR / "icon_classifier.onnx"
if not CLASSIFIER_ONNX_PATH.exists():
    CLASSIFIER_ONNX_PATH = CURRENT_DIR / "models/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 引擎初始化成功！")