""" 腾讯天御点选验证码 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)