feat: 特殊车牌生成管线 + C++/多进程/色温校正
This commit is contained in:
kirl committed
1 parent
3808b9f035
commit
3101992f51
29 files changed
+2336
-80
No files matched your search
@@ -0,0 +1,131 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
L通道车牌数据集加载器
|
||||
=====================
|
||||
Domain A: 合成车牌 (140×440) → 提取 L 通道
|
||||
Domain B: 真实车牌 (140×440) → 提取 L 通道
|
||||
"""
|
||||
|
||||
import os
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
class LChannelPlateDataset(Dataset):
|
||||
"""
|
||||
车牌 L 通道数据集。
|
||||
|
||||
Args:
|
||||
root: 图片根目录
|
||||
img_size: 统一输出尺寸 (W, H)
|
||||
is_real: True=真实车牌域, False=合成车牌域
|
||||
"""
|
||||
|
||||
def __init__(self, root: str, img_size=(440, 140), is_real=False):
|
||||
self.root = Path(root)
|
||||
self.img_size = img_size
|
||||
self.is_real = is_real
|
||||
|
||||
self.paths = []
|
||||
if self.root.is_dir():
|
||||
for ext in ('*.jpg', '*.jpeg', '*.png'):
|
||||
self.paths.extend(self.root.glob(ext))
|
||||
|
||||
def __len__(self):
|
||||
return len(self.paths)
|
||||
|
||||
def __getitem__(self, idx) -> torch.Tensor:
|
||||
path = self.paths[idx]
|
||||
img = cv2.imread(str(path))
|
||||
if img is None:
|
||||
return self.__getitem__((idx + 1) % len(self))
|
||||
|
||||
# 统一尺寸
|
||||
img = cv2.resize(img, self.img_size, interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# 提取 L 通道: BGR → LAB → L
|
||||
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
|
||||
l_channel = lab[:, :, 0].astype(np.float32) # (H, W) [0, 255]
|
||||
|
||||
# 归一化到 [-1, 1]
|
||||
l_channel = (l_channel / 127.5) - 1.0
|
||||
|
||||
return torch.from_numpy(l_channel).unsqueeze(0) # (1, H, W)
|
||||
|
||||
|
||||
def make_loader(root: str, batch_size=4, img_size=(440, 140),
|
||||
is_real=False, shuffle=True, num_workers=2):
|
||||
"""创建 DataLoader"""
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
ds = LChannelPlateDataset(root, img_size, is_real)
|
||||
return DataLoader(ds, batch_size=batch_size, shuffle=shuffle,
|
||||
num_workers=num_workers, pin_memory=True,
|
||||
drop_last=True)
|
||||
|
||||
|
||||
class LChannelInference:
|
||||
"""
|
||||
L通道 CUT 推理封装：合成车牌 → 真实风格车牌。
|
||||
|
||||
用法:
|
||||
infer = LChannelInference('checkpoint.pth')
|
||||
result_bgr = infer(synthetic_plate_bgr)
|
||||
"""
|
||||
|
||||
def __init__(self, checkpoint_path: str, device=None, n_res=6):
|
||||
from models import LChannelGenerator, init_weights
|
||||
|
||||
self.device = device or torch.device(
|
||||
"cuda" if torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
|
||||
self.G = LChannelGenerator(n_res=n_res).to(self.device)
|
||||
ckpt = torch.load(checkpoint_path, map_location=self.device)
|
||||
self.G.load_state_dict(ckpt.get('G_state', ckpt))
|
||||
self.G.eval()
|
||||
|
||||
# 期望输入尺寸
|
||||
self.img_size = (440, 140)
|
||||
|
||||
def __call__(self, plate_bgr: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Args:
|
||||
plate_bgr: (H, W, 3) 合成车牌 BGR
|
||||
|
||||
Returns:
|
||||
真实风格车牌 BGR (H, W, 3)
|
||||
"""
|
||||
h, w = plate_bgr.shape[:2]
|
||||
img = cv2.resize(plate_bgr, self.img_size, interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# BGR → LAB
|
||||
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB).astype(np.float32)
|
||||
|
||||
# L → tensor → G → tensor → L
|
||||
l_tensor = torch.from_numpy(lab[:, :, 0] / 127.5 - 1.0)
|
||||
l_tensor = l_tensor.unsqueeze(0).unsqueeze(0).to(self.device) # (1,1,H,W)
|
||||
|
||||
with torch.no_grad():
|
||||
l_out = self.G(l_tensor).squeeze().cpu().numpy() # (H, W)
|
||||
|
||||
# 还原 L 到 [0, 255]
|
||||
l_out = (l_out + 1.0) * 127.5
|
||||
l_out = np.clip(l_out, 0, 255)
|
||||
|
||||
lab[:, :, 0] = l_out
|
||||
# a, b 通道不变 → 保护颜色
|
||||
|
||||
result = cv2.cvtColor(lab.astype(np.uint8), cv2.COLOR_LAB2BGR)
|
||||
|
||||
# 还原原始尺寸
|
||||
if (h, w) != (self.img_size[1], self.img_size[0]):
|
||||
result = cv2.resize(result, (w, h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,217 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
CUT L通道 模型定义
|
||||
==================
|
||||
轻量级 Generator + PatchGAN Discriminator + PatchNCE 对比损失。
|
||||
仅处理单通道 L (亮度) 图像，输入尺寸 140×440。
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Generator: UNet 风格 encoder-decoder
|
||||
# ============================================================
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
"""残差块"""
|
||||
def __init__(self, ch):
|
||||
super().__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.ReflectionPad2d(1),
|
||||
nn.Conv2d(ch, ch, 3),
|
||||
nn.InstanceNorm2d(ch),
|
||||
nn.ReLU(True),
|
||||
nn.ReflectionPad2d(1),
|
||||
nn.Conv2d(ch, ch, 3),
|
||||
nn.InstanceNorm2d(ch),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return x + self.conv(x)
|
||||
|
||||
|
||||
class LChannelGenerator(nn.Module):
|
||||
"""
|
||||
L通道生成器: 1ch → 1ch
|
||||
输入: (B, 1, 140, 440) 合成车牌L通道
|
||||
输出: (B, 1, 140, 440) 真实风格L通道
|
||||
"""
|
||||
def __init__(self, n_res=6):
|
||||
super().__init__()
|
||||
# Encoder
|
||||
self.enc1 = self._conv_block(1, 64, 7, 1, 3) # 140×440
|
||||
self.enc2 = self._conv_block(64, 128, 3, 2, 1) # 70×220
|
||||
self.enc3 = self._conv_block(128, 256, 3, 2, 1) # 35×110
|
||||
|
||||
# Bottleneck
|
||||
self.res_blocks = nn.Sequential(*[ResBlock(256) for _ in range(n_res)])
|
||||
|
||||
# Decoder
|
||||
self.dec2 = self._upconv_block(256, 128) # 70×220
|
||||
self.dec1 = self._upconv_block(128, 64) # 140×440
|
||||
self.out = nn.Sequential(
|
||||
nn.ReflectionPad2d(3),
|
||||
nn.Conv2d(64, 1, 7),
|
||||
nn.Tanh(),
|
||||
)
|
||||
|
||||
def _conv_block(self, ci, co, k, s, p):
|
||||
return nn.Sequential(
|
||||
nn.ReflectionPad2d(p),
|
||||
nn.Conv2d(ci, co, k, s),
|
||||
nn.InstanceNorm2d(co),
|
||||
nn.ReLU(True),
|
||||
)
|
||||
|
||||
def _upconv_block(self, ci, co):
|
||||
return nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
|
||||
nn.ReflectionPad2d(1),
|
||||
nn.Conv2d(ci, co, 3),
|
||||
nn.InstanceNorm2d(co),
|
||||
nn.ReLU(True),
|
||||
)
|
||||
|
||||
def forward(self, x, return_features=False):
|
||||
"""
|
||||
Args:
|
||||
x: (B, 1, H, W) 归一化到 [-1, 1]
|
||||
return_features: 返回中间特征用于 PatchNCE loss
|
||||
"""
|
||||
features = []
|
||||
|
||||
e1 = self.enc1(x) # (B, 64, 140, 440)
|
||||
features.append(e1)
|
||||
|
||||
e2 = self.enc2(e1) # (B, 128, 70, 220)
|
||||
features.append(e2)
|
||||
|
||||
e3 = self.enc3(e2) # (B, 256, 35, 110)
|
||||
features.append(e3)
|
||||
|
||||
b = self.res_blocks(e3) # (B, 256, 35, 110)
|
||||
|
||||
d2 = self.dec2(b) # (B, 128, 70, 220)
|
||||
d1 = self.dec1(d2) # (B, 64, 140, 440)
|
||||
|
||||
out = self.out(d1) # (B, 1, 140, 440)
|
||||
|
||||
if return_features:
|
||||
return out, features
|
||||
return out
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Discriminator: PatchGAN 70×70
|
||||
# ============================================================
|
||||
|
||||
class PatchGANDiscriminator(nn.Module):
|
||||
"""
|
||||
PatchGAN 判别器，70×70 感受野。
|
||||
输入: (B, 1, 140, 440) L通道
|
||||
输出: (B, 1, H/16, W/16) 真假概率图
|
||||
"""
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.model = nn.Sequential(
|
||||
nn.Conv2d(1, 64, 4, 2, 1), # 70×220
|
||||
nn.LeakyReLU(0.2, True),
|
||||
nn.Conv2d(64, 128, 4, 2, 1), # 35×110
|
||||
nn.InstanceNorm2d(128),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
nn.Conv2d(128, 256, 4, 2, 1), # 17×55
|
||||
nn.InstanceNorm2d(256),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
nn.Conv2d(256, 512, 4, 1, 1), # 16×54
|
||||
nn.InstanceNorm2d(512),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
nn.Conv2d(512, 1, 4, 1, 1), # 15×53
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.model(x)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# PatchNCE 对比学习
|
||||
# ============================================================
|
||||
|
||||
class PatchNCELoss(nn.Module):
|
||||
"""
|
||||
PatchNCE 对比损失。
|
||||
对编码器的每一层特征图做 patch-wise 对比学习。
|
||||
"""
|
||||
def __init__(self, nce_layers=(0, 1, 2), nce_t=0.07, n_patches=256):
|
||||
super().__init__()
|
||||
self.nce_layers = nce_layers
|
||||
self.nce_t = nce_t
|
||||
self.n_patches = n_patches
|
||||
self.cross_entropy = nn.CrossEntropyLoss()
|
||||
|
||||
def forward(self, feat_q, feat_k):
|
||||
"""
|
||||
Args:
|
||||
feat_q: 生成图编码特征列表 [(B,C,H,W), ...]
|
||||
feat_k: 原图编码特征列表 [(B,C,H,W), ...]
|
||||
Returns:
|
||||
loss: scalar
|
||||
"""
|
||||
total_loss = 0.0
|
||||
n_layers = 0
|
||||
|
||||
for layer_idx in self.nce_layers:
|
||||
if layer_idx >= len(feat_q) or layer_idx >= len(feat_k):
|
||||
continue
|
||||
|
||||
fq = feat_q[layer_idx] # (B, C, H, W)
|
||||
fk = feat_k[layer_idx] # (B, C, H, W)
|
||||
|
||||
B, C, H, W = fq.shape
|
||||
n_patches = min(self.n_patches, H * W)
|
||||
|
||||
# 随机采样 patch 位置
|
||||
idx_h = torch.randint(0, H, (n_patches,), device=fq.device)
|
||||
idx_w = torch.randint(0, W, (n_patches,), device=fq.device)
|
||||
|
||||
# query patches
|
||||
q = fq[:, :, idx_h, idx_w] # (B, C, N)
|
||||
q = q.permute(0, 2, 1).reshape(B * n_patches, C) # (B*N, C)
|
||||
|
||||
# key patches (positive pairs)
|
||||
k = fk[:, :, idx_h, idx_w] # (B, C, N)
|
||||
k = k.permute(0, 2, 1).reshape(B * n_patches, C) # (B*N, C)
|
||||
|
||||
# 负样本：同一 batch 内其他图片的同一位置
|
||||
q_norm = F.normalize(q, dim=1) # (B*N, C)
|
||||
k_norm = F.normalize(k, dim=1) # (B*N, C)
|
||||
|
||||
# 计算相似度矩阵
|
||||
logits = torch.mm(q_norm, k_norm.T) / self.nce_t # (B*N, B*N)
|
||||
labels = torch.arange(B * n_patches, device=q.device)
|
||||
|
||||
loss = self.cross_entropy(logits, labels)
|
||||
total_loss += loss
|
||||
n_layers += 1
|
||||
|
||||
return total_loss / max(n_layers, 1)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 模型工具
|
||||
# ============================================================
|
||||
|
||||
def init_weights(model):
|
||||
"""Xavier 初始化"""
|
||||
for m in model.modules():
|
||||
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
|
||||
nn.init.xavier_normal_(m.weight, 0.02)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
|
||||
def count_params(model):
|
||||
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
@@ -0,0 +1,6 @@
|
||||
torch>=1.12
|
||||
torchvision
|
||||
opencv-python>=4.5
|
||||
numpy
|
||||
tqdm
|
||||
tensorboard
|
||||
@@ -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()
|
||||
Reference in new issue
Block a user