229 lines
7.7 KiB
Python
229 lines
7.7 KiB
Python
"""
|
|
图标孪生神经网络 (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)
|