升级半自动工作台
This commit is contained in:
yml2213
@@ -0,0 +1,196 @@
|
||||
"""
|
||||
腾讯天御点选验证码工业级端到端识别求解器 (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 itertools
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple
|
||||
|
||||
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_path: Path) -> List[int]:
|
||||
"""
|
||||
自动识别并解析提示图中 3 个目标图标的真实类别 ID (0~11)
|
||||
"""
|
||||
crops = slice_prompt_image(prompt_img_path)
|
||||
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
|
||||
|
||||
def _extract_candidates(self, bg_img_path: Path) -> List[Tuple[int, int, Image.Image]]:
|
||||
"""从背景图中提取所有去重后的候选图标 (cx, cy, patch)"""
|
||||
results = self.detector(str(bg_img_path), conf=0.03, verbose=False)
|
||||
boxes = results[0].boxes.xyxy.cpu().numpy()
|
||||
|
||||
bg_pil = Image.open(bg_img_path).convert("RGB")
|
||||
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_path: str | Path, prompt_img_path: str | Path) -> List[Tuple[int, int]]:
|
||||
"""
|
||||
端到端求解验证码点选坐标
|
||||
返回 3 个目标的像素坐标: [(x1, y1), (x2, y2), (x3, y3)]
|
||||
"""
|
||||
bg_path = Path(bg_img_path)
|
||||
prompt_path = Path(prompt_img_path)
|
||||
|
||||
# 1. 确定提示图的目标序列 [T1, T2, T3]
|
||||
target_classes = self.parse_prompt_targets(prompt_path)
|
||||
if len(target_classes) < 3:
|
||||
logger.warning("提示图未成功切出 3 个目标")
|
||||
return []
|
||||
|
||||
# 2. 定位背景图所有候选图标
|
||||
candidates = self._extract_candidates(bg_path)
|
||||
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}")
|
||||
Reference in New Issue
Block a user