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,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)
|
||||
Reference in new issue
Block a user