feat(onnx): 全面迁移至 ONNX Runtime，彻底移除 PyTorch 臃肿依赖
- 将 yolo_detector.pt 和 icon_classifier.pt 成功导出为 ONNX 格式 - 重构 solver.py 使用 onnxruntime + OpenCV cv2.dnn.NMSBoxes 纯轻量推理 - 彻底移除青龙对 torch, torchvision, ultralytics 等近 1.5GB 臃肿依赖 - 更新 dist/rainyun_qinglong 部署包与 README 文档
This commit is contained in:
yml2213
+109
-113
@@ -1,10 +1,11 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
腾讯天御点选验证码工业级端到端识别求解器 (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)]
|
||||
腾讯天御点选验证码工业级纯 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
|
||||
@@ -14,98 +15,74 @@ import sys
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Union
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
CURRENT_DIR = Path(__file__).resolve().parent
|
||||
PROJECT_ROOT = CURRENT_DIR.parent
|
||||
if str(PROJECT_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(PROJECT_ROOT))
|
||||
if str(CURRENT_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(CURRENT_DIR))
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.models as models
|
||||
import torchvision.transforms as transforms
|
||||
import onnxruntime as ort
|
||||
from PIL import Image
|
||||
from ultralytics import YOLO
|
||||
|
||||
from tcaptcha.slicer import slice_prompt_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")
|
||||
|
||||
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"
|
||||
# 权重路径解析 (优先 models/ 目录，兼容当前目录)
|
||||
YOLO_ONNX_PATH = CURRENT_DIR / "models/yolo_detector.onnx"
|
||||
if not YOLO_ONNX_PATH.exists():
|
||||
YOLO_ONNX_PATH = CURRENT_DIR / "yolo_detector.onnx"
|
||||
|
||||
NUM_ICON_CLASSES = 12
|
||||
TOTAL_CLASSES = 13
|
||||
CLASSIFIER_ONNX_PATH = CURRENT_DIR / "models/icon_classifier.onnx"
|
||||
if not CLASSIFIER_ONNX_PATH.exists():
|
||||
CLASSIFIER_ONNX_PATH = CURRENT_DIR / "icon_classifier.onnx"
|
||||
|
||||
|
||||
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 目标检测器 (定位图标候选位置)
|
||||
def __init__(self, yolo_path: Path = YOLO_ONNX_PATH, classifier_path: Path = CLASSIFIER_ONNX_PATH):
|
||||
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 背景干扰)
|
||||
raise FileNotFoundError(f"YOLO ONNX 权重文件不存在: {yolo_path}")
|
||||
if not classifier_path.exists():
|
||||
raise FileNotFoundError(f"Icon classifier weights not found: {classifier_path}")
|
||||
raise FileNotFoundError(f"Classifier ONNX 权重文件不存在: {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()
|
||||
# 配置轻量级 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.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]),
|
||||
])
|
||||
self.yolo_sess = ort.InferenceSession(str(yolo_path), sess_options, providers=["CPUExecutionProvider"])
|
||||
self.cls_sess = ort.InferenceSession(str(classifier_path), sess_options, providers=["CPUExecutionProvider"])
|
||||
|
||||
# 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))
|
||||
@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: str | Path | bytes | Image.Image) -> List[int]:
|
||||
"""
|
||||
自动识别并解析提示图中 3 个目标图标的真实类别 ID (0~11)
|
||||
支持干净模板匹配与 ResNet18 分类器自动双模降级
|
||||
"""
|
||||
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 []
|
||||
|
||||
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)
|
||||
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]
|
||||
|
||||
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)"""
|
||||
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):
|
||||
@@ -113,19 +90,44 @@ class CaptchaSolver:
|
||||
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()
|
||||
orig_w, orig_h = bg_pil.size
|
||||
img_np = np.array(bg_pil)
|
||||
|
||||
w_img, h_img = bg_pil.size
|
||||
# 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)))
|
||||
|
||||
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 的框合并为一个中心点
|
||||
# 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):
|
||||
@@ -133,16 +135,20 @@ class CaptchaSolver:
|
||||
|
||||
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))
|
||||
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: str | Path | bytes | Image.Image, prompt_img: str | Path | bytes | Image.Image) -> List[Tuple[int, int]]:
|
||||
def solve(
|
||||
self,
|
||||
bg_img: Union[str, Path, bytes, Image.Image],
|
||||
prompt_img: Union[str, Path, bytes, Image.Image],
|
||||
) -> List[Tuple[int, int]]:
|
||||
"""
|
||||
端到端求解验证码点选坐标
|
||||
返回 3 个目标的像素坐标: [(x1, y1), (x2, y2), (x3, y3)]
|
||||
端到端纯 ONNX 求解验证码点选坐标
|
||||
返回 3 个目标的像素坐标序列: [(x1, y1), (x2, y2), (x3, y3)]
|
||||
"""
|
||||
# 1. 确定提示图的目标序列 [T1, T2, T3]
|
||||
target_classes = self.parse_prompt_targets(prompt_img)
|
||||
@@ -156,15 +162,14 @@ class CaptchaSolver:
|
||||
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]
|
||||
# 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))
|
||||
# icon_score = P(candidate = Target_Class) * (1 - P(candidate = Distractor))
|
||||
bg_probs = probs[:, 12] # [N_candidates]
|
||||
# 4. 全局后验概率优化 (结合背景干扰项抑制因子 1 - P(distractor))
|
||||
bg_probs = probs[:, 12]
|
||||
clean_probs = probs[:, :12] * (1.0 - bg_probs[:, None])
|
||||
|
||||
num_cands = len(candidates)
|
||||
@@ -173,40 +178,31 @@ class CaptchaSolver:
|
||||
|
||||
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]]
|
||||
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 = []
|
||||
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)
|
||||
comb_list.append(best_c)
|
||||
used.add(best_c)
|
||||
best_comb = tuple(best_comb_list)
|
||||
best_comb = tuple(comb_list)
|
||||
|
||||
# 5. 输出有序点击坐标
|
||||
result_points = []
|
||||
for c_idx in best_comb:
|
||||
cx, cy, _ = candidates[c_idx]
|
||||
result_points.append((cx, cy))
|
||||
|
||||
return result_points
|
||||
return [(candidates[idx][0], candidates[idx][1]) for idx in best_comb]
|
||||
|
||||
|
||||
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}")
|
||||
print("✓ CaptchaSolver ONNX 引擎初始化成功！")
|
||||
|
||||
Reference in New Issue
Block a user