PokemonCardSearch/code/train_dinov2_layer3_bottom.py

# -*- coding: utf-8 -*-
"""
train_dinov2_layer3_bottom.py - Layer3 背景/版本消歧模型训练脚本

对齐 A4（train_dinov2_pokemonCN2_.py）的训练框架，但做了以下改动（详见
D:/顾工交接/wzj/data_all/卡牌区域/球闪/Layer3专用背景分类模型训练方案.md）：

1. backbone: facebook/dinov2-base -> facebook/dinov2-large
2. 冻结策略: 冻结前6/12层 -> 冻结前18/24层（大模型+小数据集，更保守）
3. 输入: 整卡392x392 -> 下半区392(宽)x196(高)（只训背景/版本消歧，不训物种识别）
4. 困难负样本来源: 文件夹名解析"卡名,卡号" -> 直接读 train_groups.json/val_groups.json
   （来自 find_confusable_pairs_v2.py 实测挖矿结果，group内的card_id互相是DINOv2现在会混淆的对象）
5. 数据增强: 去掉 RandomGlareA / RandomSunFlare（那是让模型对反光免疫，方向相反）
   也去掉 MotionBlur/GaussianBlur/CoarseDropout（会抹掉反光纹理细节）
   只保留：轻度亮度对比度、轻度旋转、轻度透视、轻度噪声/压缩伪影
6. 新增真正的held-out验证集（val_groups.json，46组/116张），而非训练集抽样自测

用法（pytorch环境，GPU）：
  CUDA_VISIBLE_DEVICES=0 ~/miniconda3/envs/pytorch/bin/python train_dinov2_layer3_bottom.py
"""
import os

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

import json
import random
from collections import defaultdict

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import albumentations as A
from albumentations.pytorch import ToTensorV2
from torch.utils.data import BatchSampler, DataLoader, Dataset
from tqdm import tqdm
from transformers import Dinov2Model

# ================= 配置区域 =================
DATA_ROOT = r"/home/user/顾工交接/wzj/layer3_data"
IMG_DIR = os.path.join(DATA_ROOT, "images")
TRAIN_GROUPS_JSON = os.path.join(DATA_ROOT, "train_groups.json")
VAL_GROUPS_JSON = os.path.join(DATA_ROOT, "val_groups.json")

SAVE_DIR = r"/home/user/顾工交接/wzj/layer3_bg_model_output"
BEST_MODEL_PATH = os.path.join(SAVE_DIR, "best_layer3_bottom_model.pth")

HF_MODEL_ID = "facebook/dinov2-large"

IMG_HEIGHT = 196   # 下半区裁剪高度 (392 letterbox 的下50%)
IMG_WIDTH = 392    # 下半区裁剪宽度 (全宽)
FREEZE_BLOCKS = 18  # dinov2-large 共24层，冻结前18层，只训后6层（比A4的12层冻一半更保守）

BATCH_SIZE = 8
EPOCHS = 150
LEARNING_RATE = 1e-5   # 比A4的2e-5更低：大模型+小数据集，降低过拟合/震荡风险
WEIGHT_DECAY = 1e-4
TEMPERATURE = 0.07
EVAL_INTERVAL = 1

HARD_NEGATIVE_PROB = 0.5     # 数据集本身几乎全是困难负样本组，比例可以比A4(0.35)更高
HARD_GROUP_SIZE = 6

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# ================= 数据增强 =================
# 注意：不用 RandomGlareA / RandomSunFlare（那是A4用来让模型对反光"免疫"的，
# Layer3 的目标恰好相反——要对反光纹理保持敏感），也不用 MotionBlur/GaussianBlur/
# CoarseDropout（会抹掉我们想学习的反光细节纹理）。
transform_clean = A.Compose([
    A.Resize(IMG_HEIGHT, IMG_WIDTH, interpolation=cv2.INTER_CUBIC),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2(),
])

transform_aug = A.Compose([
    A.Resize(IMG_HEIGHT, IMG_WIDTH, interpolation=cv2.INTER_CUBIC),
    A.Perspective(scale=(0.02, 0.05), p=0.4),
    A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.5),
    A.RandomGamma(gamma_limit=(85, 115), p=0.4),
    A.GaussNoise(std_range=(0.05, 0.15), p=0.3),        # 模拟相机ISO噪点
    A.ImageCompression(quality_range=(60, 95), p=0.3),   # 模拟JPEG压缩伪影
    A.SafeRotate(limit=5, border_mode=cv2.BORDER_CONSTANT, fill=0, p=0.4),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2(),
])


# ================= 1. 数据集：直接读挖矿产出的分组JSON =================
class MiningGroupDataset(Dataset):
    """
    groups_json: [{"group_id": int, "card_ids": [...], ...}, ...]
    每个 card_id 对应 IMG_DIR/{card_id}.jpg 一张背景裁剪图。
    label = 该 card_id 的唯一序号（用于检索评估，判断"能不能认回自己"）
    group_key = group_id（用于 HardNegativeBatchSampler，组内互为困难负样本）
    """
    def __init__(self, groups_json_path, img_dir, is_train):
        with open(groups_json_path, "r", encoding="utf-8") as f:
            groups = json.load(f)

        self.samples = []  # (path, label_idx, group_key)
        label_idx = 0
        for g in groups:
            group_key = g["group_id"]
            for cid in g["card_ids"]:
                path = os.path.join(img_dir, f"{cid}.jpg")
                if not os.path.exists(path):
                    continue
                self.samples.append((path, label_idx, group_key))
                label_idx += 1

        self.group_to_indices = defaultdict(list)
        for idx, (_, _, group_key) in enumerate(self.samples):
            self.group_to_indices[group_key].append(idx)

        self.valid_hard_groups = [g for g, idxs in self.group_to_indices.items() if len(idxs) >= 2]

        tag = "train" if is_train else "val"
        print(f"[{tag}] samples={len(self.samples)}  hard_negative_groups(>=2)={len(self.valid_hard_groups)}")

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        path, label, group_key = self.samples[idx]
        img = cv2.imread(path)
        if img is None:
            img = np.zeros((IMG_HEIGHT, IMG_WIDTH, 3), dtype=np.uint8)
        else:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        img_clean = transform_clean(image=img)["image"]
        img_aug = transform_aug(image=img)["image"]
        return img_clean, img_aug, label


# ================= 2. 困难负样本 BatchSampler（原样复用 A4 逻辑）=================
class HardNegativeBatchSampler(BatchSampler):
    def __init__(self, dataset, batch_size, hard_negative_prob=0.35, hard_group_size=8, drop_last=True):
        self.dataset = dataset
        self.batch_size = batch_size
        self.hard_negative_prob = hard_negative_prob
        self.hard_group_size = min(hard_group_size, batch_size)
        self.drop_last = drop_last
        self.num_samples = len(dataset)
        self.all_indices = list(range(self.num_samples))

    def __iter__(self):
        unused = set(self.all_indices)
        while len(unused) >= self.batch_size:
            batch = []
            if len(self.dataset.valid_hard_groups) > 0 and random.random() < self.hard_negative_prob:
                candidate_groups = []
                for g in self.dataset.valid_hard_groups:
                    available = [i for i in self.dataset.group_to_indices[g] if i in unused]
                    if len(available) >= 2:
                        candidate_groups.append((g, available))
                if candidate_groups:
                    g, available = random.choice(candidate_groups)
                    take_n = min(len(available), self.hard_group_size, self.batch_size)
                    chosen_hard = random.sample(available, take_n)
                    batch.extend(chosen_hard)
                    for i in chosen_hard:
                        unused.remove(i)

            remain = self.batch_size - len(batch)
            if remain > 0:
                if len(unused) < remain:
                    if self.drop_last:
                        break
                    batch.extend(list(unused))
                    unused.clear()
                else:
                    chosen_rand = random.sample(list(unused), remain)
                    batch.extend(chosen_rand)
                    for i in chosen_rand:
                        unused.remove(i)

            if len(batch) == self.batch_size:
                random.shuffle(batch)
                yield batch

        if not self.drop_last and unused:
            yield list(unused)

    def __len__(self):
        if self.drop_last:
            return self.num_samples // self.batch_size
        return (self.num_samples + self.batch_size - 1) // self.batch_size


# ================= 3. 模型：DINOv2-large，冻结前18层 =================
class Dinov2Layer3Model(nn.Module):
    def __init__(self, freeze_blocks=FREEZE_BLOCKS):
        super().__init__()
        print(f"Loading {HF_MODEL_ID} backbone ...")
        self.backbone = Dinov2Model.from_pretrained(HF_MODEL_ID)

        for param in self.backbone.parameters():
            param.requires_grad = False

        total_layers = len(self.backbone.encoder.layer)
        print(f"Backbone total layers: {total_layers}, freezing first {freeze_blocks}")
        for i in range(freeze_blocks, total_layers):
            for param in self.backbone.encoder.layer[i].parameters():
                param.requires_grad = True
        for param in self.backbone.layernorm.parameters():
            param.requires_grad = True

    def forward(self, x):
        # 输入非正方形(392x196)，需要开启位置编码插值
        outputs = self.backbone(x, interpolate_pos_encoding=True)
        cls_token = outputs.last_hidden_state[:, 0, :]
        return F.normalize(cls_token, p=2, dim=1)


# ================= 4. InfoNCE 对比损失（原样复用 A4）=================
def contrastive_loss(feat_clean, feat_aug, temperature=TEMPERATURE):
    batch_size = feat_clean.size(0)
    logits = torch.matmul(feat_clean, feat_aug.T) / temperature
    labels = torch.arange(batch_size).to(feat_clean.device)
    loss_1 = F.cross_entropy(logits, labels)
    loss_2 = F.cross_entropy(logits.T, labels)
    return (loss_1 + loss_2) / 2.0


# ================= 5. 检索式评估（在真正held-out的val集上）=================
def evaluate(model, val_dataset):
    model.eval()
    loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)

    gallery_feats, gallery_labels = [], []
    query_feats, query_labels = [], []

    with torch.no_grad():
        for imgs_c, imgs_a, lbls in loader:
            imgs_c, imgs_a = imgs_c.to(DEVICE), imgs_a.to(DEVICE)
            with torch.amp.autocast("cuda"):
                f_c = model(imgs_c)
                f_a = model(imgs_a)
            gallery_feats.append(f_c.cpu())
            gallery_labels.append(lbls)
            query_feats.append(f_a.cpu())
            query_labels.append(lbls)

    gallery_feats = torch.cat(gallery_feats, dim=0)
    gallery_labels = torch.cat(gallery_labels, dim=0)
    query_feats = torch.cat(query_feats, dim=0)
    query_labels = torch.cat(query_labels, dim=0)

    sim_matrix = torch.matmul(query_feats, gallery_feats.T)
    k = min(5, gallery_feats.size(0))
    topk_idx = sim_matrix.topk(k, dim=1).indices

    correct_1, correct_5, total = 0, 0, query_labels.size(0)
    for idx in range(total):
        q_label = query_labels[idx]
        retrieved = gallery_labels[topk_idx[idx]]
        if q_label == retrieved[0]:
            correct_1 += 1
        if q_label in retrieved:
            correct_5 += 1

    return correct_1 / total, correct_5 / total


# ================= 6. 主流程 =================
def main():
    os.makedirs(SAVE_DIR, exist_ok=True)

    train_dataset = MiningGroupDataset(TRAIN_GROUPS_JSON, IMG_DIR, is_train=True)
    val_dataset = MiningGroupDataset(VAL_GROUPS_JSON, IMG_DIR, is_train=False)

    batch_sampler = HardNegativeBatchSampler(
        dataset=train_dataset,
        batch_size=BATCH_SIZE,
        hard_negative_prob=HARD_NEGATIVE_PROB,
        hard_group_size=HARD_GROUP_SIZE,
        drop_last=True,
    )
    train_loader = DataLoader(train_dataset, batch_sampler=batch_sampler, num_workers=8, pin_memory=True)

    model = Dinov2Layer3Model().to(DEVICE)
    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)

    optimizer = optim.AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY,
    )
    scaler = torch.amp.GradScaler("cuda")

    best_top1_acc = 0.0
    print("开始 Layer3 (下半区/背景版本消歧) 对比学习训练 ...")

    for epoch in range(EPOCHS):
        model.train()
        pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{EPOCHS} [Train]")
        running_loss = 0.0

        for i, (imgs_clean, imgs_aug, _) in enumerate(pbar):
            imgs_clean = imgs_clean.to(DEVICE)
            imgs_aug = imgs_aug.to(DEVICE)

            optimizer.zero_grad()
            with torch.amp.autocast("cuda"):
                feat_clean = model(imgs_clean)
                feat_aug = model(imgs_aug)
                loss = contrastive_loss(feat_clean, feat_aug, TEMPERATURE)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            running_loss += loss.item()
            pbar.set_postfix({"InfoNCE Loss": f"{running_loss / (i + 1):.4f}"})

        if (epoch + 1) % EVAL_INTERVAL == 0:
            top1_acc, top5_acc = evaluate(model, val_dataset)
            print(f"[Epoch {epoch + 1} Eval(held-out val)] Retrieval Top-1: {top1_acc:.2%} | Top-5: {top5_acc:.2%}")

            if top1_acc > best_top1_acc:
                best_top1_acc = top1_acc
                raw_model = model.module if hasattr(model, "module") else model
                torch.save(raw_model.state_dict(), BEST_MODEL_PATH)
                print(f"保存新的最佳模型! Top-1: {top1_acc:.2%} -> {BEST_MODEL_PATH}")


if __name__ == "__main__":
    main()