升级半自动工作台

2026-09-11 22:56:03 +08:00
parent f1c664f622
commit 4d6cded0b0
13 changed files with 3269 additions and 231 deletions
@@ -0,0 +1,241 @@
+"""
+腾讯天御点选验证码 12 类精确图标分类器训练脚本 (ResNet18 + 迁移学习 + 真实与合成增强)
+类别定义: 0~11 为 12 种独立图形图标，12 为背景干扰类 (共 13 类)
+"""
+
+import json
+import logging
+import random
+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
+from PIL import Image
+from torch.utils.data import DataLoader, Dataset
+import torchvision.models as models
+import torchvision.transforms as transforms
+from tcaptcha.slicer import slice_prompt_image
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("train_classifier")
+
+MODEL_SAVE_PATH = PROJECT_ROOT / "tcaptcha/models/icon_classifier.pt"
+TEMPLATES_DIR = PROJECT_ROOT / "tcaptcha/data/templates"
+LABELS_FILE = PROJECT_ROOT / "tcaptcha/data/labels.jsonl"
+RAW_DIR = PROJECT_ROOT / "tcaptcha/data/raw"
+CROPS_DIR = PROJECT_ROOT / "tcaptcha/data/crops"
+
+NUM_ICON_CLASSES = 12
+TOTAL_CLASSES = 13  # 0~11 图标类 + 12 干扰背景类
+
+
+def load_templates():
+    templates = []
+    for cid in range(NUM_ICON_CLASSES):
+        p = TEMPLATES_DIR / f"class_{cid}.png"
+        img = Image.open(p).convert("RGBA")
+        arr = np.array(img.convert("L"))
+        bg = np.median(arr)
+        bin_t = cv2.resize(((arr < bg - 20).astype(np.uint8) * 255), (48, 48), interpolation=cv2.INTER_AREA)
+        templates.append((cid, img, bin_t))
+    return templates
+
+
+class CaptchaIconDataset(Dataset):
+    def __init__(self, samples: List[Tuple[Image.Image, int]], transform=None):
+        self.samples = samples
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        img, label = self.samples[idx]
+        if self.transform:
+            img = self.transform(img)
+        return img, label
+
+
+def generate_synthetic_samples(templates, bg_images, num_per_class: int = 150):
+    """通过将无旋转纯净模板以随机角度、随机透明度印在真实背景切片上，合成大量高质量带标签数据"""
+    synthetic_samples = []
+    for cid, t_img, _ in templates:
+        for _ in range(num_per_class):
+            bg = random.choice(bg_images)
+            w_bg, h_bg = bg.size
+            rx = random.randint(10, max(11, w_bg - 70))
+            ry = random.randint(10, max(11, h_bg - 70))
+            patch = bg.crop((rx, ry, rx + 64, ry + 64)).copy()
+
+            # 随机旋转与缩放
+            ang = random.uniform(-45, 45)
+            rot_t = t_img.rotate(ang, expand=True, resample=Image.Resampling.BICUBIC)
+            target_size = random.randint(40, 52)
+            rot_t = rot_t.resize((target_size, target_size), Image.Resampling.LANCZOS)
+
+            # 掩码
+            t_arr = np.array(rot_t.convert("L"))
+            alpha_val = random.randint(160, 240)
+            mask = ((t_arr < 135).astype(np.uint8) * alpha_val)
+
+            t_stamp = Image.new("RGBA", (64, 64), (0, 0, 0, 0))
+            offset_x = (64 - target_size) // 2 + random.randint(-3, 3)
+            offset_y = (64 - target_size) // 2 + random.randint(-3, 3)
+            # 深色线条
+            color_val = random.randint(15, 60)
+            icon_color = Image.new("RGBA", (target_size, target_size), (color_val, color_val, color_val, 255))
+            t_stamp.paste(icon_color, (offset_x, offset_y), mask=Image.fromarray(mask))
+
+            blended = Image.alpha_composite(patch, t_stamp).convert("RGB")
+            synthetic_samples.append((blended, cid))
+
+    # 随机采样纯背景负样本 (Class 12: 背景干扰，与各图标类平衡)
+    for _ in range(num_per_class):
+        bg = random.choice(bg_images)
+        w_bg, h_bg = bg.size
+        rx = random.randint(10, max(11, w_bg - 70))
+        ry = random.randint(10, max(11, h_bg - 70))
+        patch = bg.crop((rx, ry, rx + 64, ry + 64)).convert("RGB")
+        synthetic_samples.append((patch, 12))
+
+    return synthetic_samples
+
+
+def build_real_samples(templates):
+    """从人工标注的真实切片中提取正样本"""
+    records = [json.loads(line) for line in open(LABELS_FILE, "r", encoding="utf-8")]
+    real_samples = []
+
+    for r in records:
+        sid = r["sample_id"]
+        crops = slice_prompt_image(RAW_DIR / f"{sid}_prompt.png")
+        for i in range(3):
+            p_c = crops[i][0]
+            arr = np.array(p_c.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[2], cv2.TM_CCOEFF_NORMED)[0, 0] for t in templates]
+            cid = int(np.argmax(scores))
+
+            t_crop_path = CROPS_DIR / f"{sid}_target_{i+1}.png"
+            if t_crop_path.exists():
+                img = Image.open(t_crop_path).convert("RGB")
+                real_samples.append((img, cid))
+
+    return real_samples
+
+
+def train(epochs: int = 15, batch_size: int = 32, lr: float = 0.0008):
+    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+    logger.info("使用加速计算设备: %s", device)
+
+    templates = load_templates()
+    logger.info("已加载 %d 个基准图标模板", len(templates))
+
+    bg_images = []
+    for sid_meta in list(RAW_DIR.glob("*_bg.png"))[:35]:
+        bg_images.append(Image.open(sid_meta).convert("RGBA"))
+
+    logger.info("正在生成 12 类合成旋转与贴图数据集...")
+    synthetic = generate_synthetic_samples(templates, bg_images, num_per_class=120)
+    real = build_real_samples(templates)
+
+    logger.info("真实标注图块: %d, 合成增强图块: %d, 总计: %d", len(real), len(synthetic), len(real) + len(synthetic))
+
+    # 全量生产级训练: 真实样本过采样 6 倍 + 合成样本
+    logger.info("采用全量样本训练生产级模型 (包含 100% 真实人工标注样本)...")
+    train_samples = real * 6 + synthetic
+    val_samples = real
+
+    random.shuffle(train_samples)
+
+    train_transform = transforms.Compose([
+        transforms.Resize((64, 64)),
+        transforms.RandomRotation(degrees=20),
+        transforms.RandomResizedCrop(64, scale=(0.85, 1.05)),
+        transforms.ColorJitter(brightness=0.2, contrast=0.2),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    val_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]),
+    ])
+
+    train_loader = DataLoader(CaptchaIconDataset(train_samples, train_transform), batch_size=batch_size, shuffle=True)
+    val_loader = DataLoader(CaptchaIconDataset(val_samples, val_transform), batch_size=batch_size, shuffle=False)
+
+    # 采用预训练 ResNet18
+    model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
+    # 修改分类头为 TOTAL_CLASSES 类 (0~11 图标类 + 12 背景干扰类)
+    model.fc = nn.Linear(model.fc.in_features, TOTAL_CLASSES)
+    model = model.to(device)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+
+    logger.info("开始训练 ResNet18 13类图标分类器 (Train: %d, Val Real: %d)...", len(train_samples), len(val_samples))
+
+    best_acc = 0.0
+    for epoch in range(1, epochs + 1):
+        model.train()
+        total_loss = 0.0
+        correct_train = 0
+        total_train = 0
+
+        for imgs, labels in train_loader:
+            imgs, labels = imgs.to(device), labels.to(device)
+            optimizer.zero_grad()
+            outputs = model(imgs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item()
+            _, preds = outputs.max(1)
+            correct_train += preds.eq(labels).sum().item()
+            total_train += labels.size(0)
+
+        scheduler.step()
+
+        # 验证
+        model.eval()
+        correct_val = 0
+        total_val = 0
+        with torch.no_grad():
+            for imgs, labels in val_loader:
+                imgs, labels = imgs.to(device), labels.to(device)
+                outputs = model(imgs)
+                _, preds = outputs.max(1)
+                correct_val += preds.eq(labels).sum().item()
+                total_val += labels.size(0)
+
+        val_acc = correct_val / total_val if total_val > 0 else 0
+        if val_acc >= best_acc or epoch == epochs:
+            best_acc = max(val_acc, best_acc)
+            MODEL_SAVE_PATH.parent.mkdir(parents=True, exist_ok=True)
+            torch.save(model.state_dict(), MODEL_SAVE_PATH)
+
+        logger.info(
+            "Epoch [%02d/%02d] Loss: %.4f | Train Acc: %.2f%% | Val Real Acc: %.2f%% (Best: %.2f%%)",
+            epoch, epochs, total_loss / len(train_loader),
+            correct_train / total_train * 100,
+            val_acc * 100, best_acc * 100
+        )
+
+    logger.info("✓ 分类模型训练完成，最佳真实验证集准确率: %.2f%%，已保存至: %s", best_acc * 100, MODEL_SAVE_PATH)
+
+
+if __name__ == "__main__":
+    train(epochs=20)