218 lines
6.7 KiB
Python
218 lines
6.7 KiB
Python
# -*- 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)
|