feat: 特殊车牌生成管线 + C++/多进程/色温校正

2026-06-21 10:32:37 +08:00
1 parent 3808b9f035
commit 3101992f51
@@ -0,0 +1,204 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+CUT L通道 训练脚本
+==================
+仅对车牌 L (亮度) 通道做纹理迁移训练，a,b 颜色通道不动。
+
+用法:
+  python train.py \
+      --domain-a ./gan_data/domain_A_synthetic/police \
+      --domain-b ./gan_data/domain_B_real \
+      --epochs 100 \
+      --batch 4 \
+      --output ./output/cut_police
+"""
+
+import os
+import sys
+import argparse
+import time
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+# 添加 tools 目录到路径
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from models import (
+    LChannelGenerator, PatchGANDiscriminator,
+    PatchNCELoss, init_weights, count_params,
+)
+from dataset import make_loader
+
+
+def train(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    # ---- 数据 ----
+    loader_a = make_loader(args.domain_a, args.batch, (args.img_w, args.img_h),
+                           is_real=False, num_workers=args.workers)
+    loader_b = make_loader(args.domain_b, args.batch, (args.img_w, args.img_h),
+                           is_real=True, num_workers=args.workers)
+
+    print(f"Domain A (合成): {len(loader_a.dataset)} 张")
+    print(f"Domain B (真实): {len(loader_b.dataset)} 张")
+    print(f"Batch size:      {args.batch}")
+    print(f"Epochs:          {args.epochs}")
+
+    # ---- 模型 ----
+    G = LChannelGenerator(n_res=args.n_res).to(device)
+    D = PatchGANDiscriminator().to(device)
+
+    init_weights(G)
+    init_weights(D)
+
+    print(f"G params: {count_params(G):,}")
+    print(f"D params: {count_params(D):,}")
+
+    # ---- 损失 & 优化器 ----
+    criterion_gan = nn.MSELoss()           # LSGAN
+    criterion_nce = PatchNCELoss(
+        nce_layers=args.nce_layers,
+        nce_t=args.nce_t,
+        n_patches=args.n_patches,
+    ).to(device)
+    criterion_idt = nn.L1Loss()
+
+    opt_G = optim.Adam(G.parameters(), lr=args.lr, betas=(0.5, 0.999))
+    opt_D = optim.Adam(D.parameters(), lr=args.lr, betas=(0.5, 0.999))
+
+    # ---- 日志 ----
+    os.makedirs(args.output, exist_ok=True)
+
+    # ---- 训练循环 ----
+    for epoch in range(1, args.epochs + 1):
+        epoch_start = time.time()
+        G_loss_total = 0.0
+        D_loss_total = 0.0
+        n_batches = 0
+
+        for real_a, real_b in zip(loader_a, loader_b):
+            real_a = real_a.to(device)      # (B, 1, H, W) 合成L
+            real_b = real_b.to(device)      # (B, 1, H, W) 真实L
+            B = real_a.size(0)
+
+            # ---- 1. 更新 D ----
+            fake_b = G(real_a).detach()
+            # G 不用返回特征
+
+            pred_real = D(real_b)
+            pred_fake = D(fake_b)
+
+            loss_D_real = criterion_gan(pred_real,
+                                         torch.ones_like(pred_real))
+            loss_D_fake = criterion_gan(pred_fake,
+                                         torch.zeros_like(pred_fake))
+            loss_D = (loss_D_real + loss_D_fake) * 0.5
+
+            opt_D.zero_grad()
+            loss_D.backward()
+            opt_D.step()
+
+            # ---- 2. 更新 G ----
+            fake_b, feat_q = G(real_a, return_features=True)
+
+            # GAN loss
+            pred_fake = D(fake_b)
+            loss_G_gan = criterion_gan(pred_fake,
+                                        torch.ones_like(pred_fake))
+
+            # PatchNCE loss (对比学习)
+            _, feat_k = G(real_b, return_features=True)
+            loss_G_nce = criterion_nce(feat_q, feat_k)
+
+            # Identity loss: 裁剪输出到输入尺寸（UNet可能多出几像素）
+            identity = G(real_a)
+            if identity.shape != real_a.shape:
+                identity = identity[:, :, :real_a.shape[2], :real_a.shape[3]]
+            loss_G_idt = criterion_idt(identity, real_a) * args.lambda_idt
+
+            loss_G = loss_G_gan + loss_G_nce * args.lambda_nce + loss_G_idt
+
+            opt_G.zero_grad()
+            loss_G.backward()
+            opt_G.step()
+
+            G_loss_total += loss_G.item()
+            D_loss_total += loss_D.item()
+            n_batches += 1
+
+        # ---- Epoch 统计 ----
+        G_avg = G_loss_total / max(n_batches, 1)
+        D_avg = D_loss_total / max(n_batches, 1)
+        elapsed = time.time() - epoch_start
+
+        print(f"Epoch {epoch:3d}/{args.epochs} | "
+              f"G: {G_avg:.4f}  D: {D_avg:.4f} | "
+              f"{elapsed:.1f}s")
+
+        pass  # 日志仅print，不使用tensorboard
+
+        # ---- 保存 ----
+        if epoch % args.save_every == 0:
+            ckpt = {
+                'epoch': epoch,
+                'G_state': G.state_dict(),
+                'D_state': D.state_dict(),
+                'opt_G': opt_G.state_dict(),
+                'opt_D': opt_D.state_dict(),
+            }
+            torch.save(ckpt, os.path.join(args.output, f'epoch_{epoch}.pth'))
+
+    # 最终保存
+    final_ckpt = {
+        'G_state': G.state_dict(),
+        'D_state': D.state_dict(),
+    }
+    torch.save(final_ckpt, os.path.join(args.output, 'final.pth'))
+    print(f"\n训练完成! 模型保存在: {args.output}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description='CUT L通道训练')
+
+    # 数据
+    parser.add_argument('--domain-a', type=str, required=True,
+                        help='合成车牌目录')
+    parser.add_argument('--domain-b', type=str, required=True,
+                        help='真实车牌目录')
+
+    # 训练
+    parser.add_argument('--epochs', type=int, default=100)
+    parser.add_argument('--batch', type=int, default=4)
+    parser.add_argument('--lr', type=float, default=2e-4)
+    parser.add_argument('--workers', type=int, default=2)
+
+    # 模型
+    parser.add_argument('--n-res', type=int, default=6,
+                        help='残差块数量')
+
+    # 损失权重
+    parser.add_argument('--lambda-nce', type=float, default=1.0)
+    parser.add_argument('--lambda-idt', type=float, default=0.5)
+
+    # PatchNCE
+    parser.add_argument('--nce-layers', type=int, nargs='+', default=[0, 1, 2])
+    parser.add_argument('--nce-t', type=float, default=0.07)
+    parser.add_argument('--n-patches', type=int, default=128)
+
+    # 保存
+    parser.add_argument('--output', type=str, default='./output/cut_train')
+    parser.add_argument('--img-w', type=int, default=220)
+    parser.add_argument('--img-h', type=int, default=70)
+    parser.add_argument('--save-every', type=int, default=20)
+
+    args = parser.parse_args()
+    train(args)
+
+
+if __name__ == '__main__':
+    main()