""" 图标孪生神经网络 (Icon Siamese Network) 用于计算提示图图标与背景候选图标之间的形状相似度（抗背景噪点、抗旋转） """ import json import logging import random from pathlib import Path from typing import List, Tuple import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torchvision.transforms as transforms from PIL import Image from torch.utils.data import DataLoader, Dataset import sys PROJECT_ROOT = Path(__file__).resolve().parent.parent if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from tcaptcha.slicer import slice_prompt_image logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") logger = logging.getLogger("siamese") MODEL_PATH = PROJECT_ROOT / "tcaptcha/models/siamese_matcher.pt" LABEL_FILE = PROJECT_ROOT / "tcaptcha/data/labels.jsonl" RAW_DIR = PROJECT_ROOT / "tcaptcha/data/raw" CROPS_DIR = PROJECT_ROOT / "tcaptcha/data/crops" class IconSiameseNet(nn.Module): """轻量级 CNN 孪生特征提取网络""" def __init__(self, emb_dim: int = 128): super().__init__() self.encoder = nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.LeakyReLU(0.1, inplace=True), nn.MaxPool2d(2, 2), # 32x32 nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.LeakyReLU(0.1, inplace=True), nn.MaxPool2d(2, 2), # 16x16 nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.LeakyReLU(0.1, inplace=True), nn.MaxPool2d(2, 2), # 8x8 nn.Conv2d(128, 256, kernel_size=3, padding=1), nn.BatchNorm2d(256), nn.LeakyReLU(0.1, inplace=True), nn.AdaptiveAvgPool2d((1, 1)), # 256x1x1 ) self.head = nn.Sequential( nn.Flatten(), nn.Linear(256, emb_dim), ) def forward_once(self, x: torch.Tensor) -> torch.Tensor: feat = self.encoder(x) emb = self.head(feat) return F.normalize(emb, p=2, dim=1) def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: e1 = self.forward_once(x1) e2 = self.forward_once(x2) # 余弦相似度 [-1, 1] 映射到 [0, 1] cosine = (e1 * e2).sum(dim=1) return (cosine + 1.0) / 2.0 class IconPairDataset(Dataset): """构建 (Prompt图标, Target图块) 正负样本对""" def __init__(self, pairs: List[Tuple[Image.Image, Image.Image, int]], transform=None): self.pairs = pairs self.transform = transform def __len__(self): return len(self.pairs) def __getitem__(self, idx): img1, img2, label = self.pairs[idx] if self.transform: img1 = self.transform(img1) img2 = self.transform(img2) return img1, img2, torch.tensor(label, dtype=torch.float32) def load_pairs_from_labels(): if not LABEL_FILE.exists(): return [], [] with open(LABEL_FILE, "r", encoding="utf-8") as f: records = [json.loads(line) for line in f if line.strip()] all_prompts = [] all_targets = [] for r in records: sid = r["sample_id"] prompt_path = RAW_DIR / f"{sid}_prompt.png" if not prompt_path.exists(): continue try: p_crops = slice_prompt_image(prompt_path) for i in range(3): t_crop_path = CROPS_DIR / f"{sid}_target_{i+1}.png" if t_crop_path.exists() and i < len(p_crops): p_img = p_crops[i][0].convert("RGB") t_img = Image.open(t_crop_path).convert("RGB") all_prompts.append(p_img) all_targets.append(t_img) except Exception as e: logger.warning("解析样本 %s 异常: %s", sid, e) return all_prompts, all_targets def build_augmented_pairs(prompts, targets, num_neg_per_pos: int = 3): pairs = [] n = len(prompts) for i in range(n): # 正样本对 (label = 1) pairs.append((prompts[i], targets[i], 1)) # 负样本对 (label = 0) for _ in range(num_neg_per_pos): neg_idx = random.choice([j for j in range(n) if j != i]) pairs.append((prompts[i], targets[neg_idx], 0)) random.shuffle(pairs) return pairs def train_siamese(epochs: int = 25, batch_size: int = 32, lr: float = 0.001): MODEL_PATH.parent.mkdir(parents=True, exist_ok=True) device = torch.device("mps" if torch.backends.mps.is_available() else "cpu") logger.info("采用训练加速设备: %s", device) prompts, targets = load_pairs_from_labels() logger.info("读取到 %d 组高质量正样本对应关系", len(prompts)) if len(prompts) < 10: logger.error("标注样本不足，无法训练孪生网络") return # 数据集划分 (80% 训练, 20% 验证) indices = list(range(len(prompts))) random.shuffle(indices) split = int(len(prompts) * 0.8) train_p = [prompts[i] for i in indices[:split]] train_t = [targets[i] for i in indices[:split]] val_p = [prompts[i] for i in indices[split:]] val_t = [targets[i] for i in indices[split:]] train_pairs = build_augmented_pairs(train_p, train_t, num_neg_per_pos=3) val_pairs = build_augmented_pairs(val_p, val_t, num_neg_per_pos=3) train_transform = transforms.Compose([ transforms.Resize((64, 64)), transforms.RandomRotation(degrees=45), # 强化抗旋转 transforms.ColorJitter(brightness=0.2, contrast=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ]) val_transform = transforms.Compose([ transforms.Resize((64, 64)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ]) train_loader = DataLoader(IconPairDataset(train_pairs, train_transform), batch_size=batch_size, shuffle=True) val_loader = DataLoader(IconPairDataset(val_pairs, val_transform), batch_size=batch_size, shuffle=False) model = IconSiameseNet(emb_dim=128).to(device) criterion = nn.BCELoss() optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) logger.info("开始训练孪生图标比对网络 (Train: %d 对, Val: %d 对)...", len(train_pairs), len(val_pairs)) best_val_acc = 0.0 for epoch in range(1, epochs + 1): model.train() train_loss = 0.0 for x1, x2, y in train_loader: x1, x2, y = x1.to(device), x2.to(device), y.to(device) optimizer.zero_grad() pred = model(x1, x2) loss = criterion(pred, y) loss.backward() optimizer.step() train_loss += loss.item() # 验证 model.eval() correct = 0 total = 0 with torch.no_grad(): for x1, x2, y in val_loader: x1, x2, y = x1.to(device), x2.to(device), y.to(device) pred = model(x1, x2) predicted_labels = (pred > 0.5).float() correct += (predicted_labels == y).sum().item() total += y.size(0) val_acc = correct / total if total > 0 else 0 if val_acc > best_val_acc: best_val_acc = val_acc torch.save(model.state_dict(), MODEL_PATH) if epoch % 5 == 0 or epoch == epochs: logger.info("Epoch [%d/%d] Loss: %.4f | Val Pair Acc: %.2f%% (Best: %.2f%%)", epoch, epochs, train_loss / len(train_loader), val_acc * 100, best_val_acc * 100) logger.info("✓ 孪生网络训练完成，最佳权重已保存至: %s", MODEL_PATH) if __name__ == "__main__": train_siamese(epochs=30)