📚 课程目录第4章基于GAN的低光照图像增强

4. 基于GAN的低光照图像增强

任务目标

构建一个条件GAN（cGAN）用于低光照图像增强，在生成器中使用U-Net结构并引入感知损失（VGG特征匹配），使模型在提升暗部亮度的同时恢复真实色彩与纹理细节。最终与直方图均衡化（HE）方法进行定性对比，验证GAN方法的增强效果。

环境依赖

Python 3.8+，PyTorch 1.10+，torchvision，OpenCV，NumPy，Matplotlib。需预训练VGG16权重（torchvision自动下载）。数据集使用任意成对低光照/正常光照图像（如LOL数据集），本案例以单张示例图像演示完整训练与推理流程。

核心实现

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as vgg
import cv2
import numpy as np
from torch.utils.data import Dataset, DataLoader

# ---------- 生成器：U-Net ----------
class UNetGen(nn.Module):
    def __init__(self):
        super().__init__()
        # 编码器（3→64→128→256）
        self.e1 = nn.Sequential(nn.Conv2d(3,64,4,2,1), nn.LeakyReLU(0.2))
        self.e2 = nn.Sequential(nn.Conv2d(64,128,4,2,1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2))
        self.e3 = nn.Sequential(nn.Conv2d(128,256,4,2,1), nn.BatchNorm2d(256), nn.LeakyReLU(0.2))
        # 解码器（256→128→64→3）
        self.d1 = nn.Sequential(nn.ConvTranspose2d(256,128,4,2,1), nn.BatchNorm2d(128), nn.ReLU())
        self.d2 = nn.Sequential(nn.ConvTranspose2d(256,64,4,2,1), nn.BatchNorm2d(64), nn.ReLU())
        self.d3 = nn.Sequential(nn.ConvTranspose2d(128,3,4,2,1), nn.Tanh())
    def forward(self, x):
        e1 = self.e1(x); e2 = self.e2(e1); e3 = self.e3(e2)
        d1 = self.d1(e3)
        d2 = self.d2(torch.cat([d1, e2], dim=1))
        d3 = self.d3(torch.cat([d2, e1], dim=1))
        return d3

# ---------- 判别器 ----------
class Disc(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(6,64,4,2,1), nn.LeakyReLU(0.2),
            nn.Conv2d(64,128,4,2,1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2),
            nn.Conv2d(128,256,4,2,1), nn.BatchNorm2d(256), nn.LeakyReLU(0.2),
            nn.Conv2d(256,1,4,1,1), nn.Sigmoid())
    def forward(self, low, high):
        return self.model(torch.cat([low, high], dim=1))

# ---------- 感知损失（VGG16特征匹配） ----------
class PerceptualLoss(nn.Module):
    def __init__(self):
        super().__init__()
        vgg16 = vgg.vgg16(pretrained=True).features[:16].eval()
        for p in vgg16.parameters(): p.requires_grad = False
        self.vgg = vgg16
    def forward(self, pred, target):
        return nn.functional.l1_loss(self.vgg(pred), self.vgg(target))

# ---------- 数据加载（单张示例） ----------
class PairDataset(Dataset):
    def __init__(self, low_path, high_path, size=256):
        self.low = cv2.cvtColor(cv2.resize(cv2.imread(low_path), (size,size)), cv2.COLOR_BGR2RGB)
        self.high = cv2.cvtColor(cv2.resize(cv2.imread(high_path), (size,size)), cv2.COLOR_BGR2RGB)
    def __len__(self): return 1
    def __getitem__(self, i):
        l = torch.from_numpy(self.low/127.5-1).permute(2,0,1).float()
        h = torch.from_numpy(self.high/127.5-1).permute(2,0,1).float()
        return l, h

# ---------- 训练配置 ----------
device = 'cuda' if torch.cuda.is_available() else 'cpu'
G = UNetGen().to(device)
D = Disc().to(device)
optG = optim.Adam(G.parameters(), lr=2e-4, betas=(0.5,0.999))
optD = optim.Adam(D.parameters(), lr=2e-4, betas=(0.5,0.999))
bce = nn.BCELoss()
percep = PerceptualLoss().to(device)
l1 = nn.L1Loss()
ds = PairDataset('low.jpg', 'high.jpg')
dl = DataLoader(ds, batch_size=1)

# ---------- 训练循环（200轮） ----------
for epoch in range(200):
    for low, high in dl:
        low, high = low.to(device), high.to(device)
        real_label = torch.ones(low.size(0),1,16,16, device=device)
        fake_label = torch.zeros(low.size(0),1,16,16, device=device)
        # 训练判别器
        fake = G(low)
        d_real = D(low, high)
        d_fake = D(low, fake.detach())
        lossD = bce(d_real, real_label) + bce(d_fake, fake_label)
        optD.zero_grad(); lossD.backward(); optD.step()
        # 训练生成器（对抗+感知+L1）
        d_fake2 = D(low, fake)
        lossG_adv = bce(d_fake2, real_label)
        lossG_per = percep(fake, high)
        lossG_l1 = l1(fake, high)
        lossG = lossG_adv + 10*lossG_per + 100*lossG_l1
        optG.zero_grad(); lossG.backward(); optG.step()
    if epoch % 50 == 0:
        print(f'Epoch {epoch} | D: {lossD.item():.4f} | G: {lossG.item():.4f}')

# ---------- 推理与对比 ----------
torch.save(G.state_dict(), 'gan_enhance.pth')
G.eval()
with torch.no_grad():
    low_img = cv2.cvtColor(cv2.resize(cv2.imread('low.jpg'), (256,256)), cv2.COLOR_BGR2RGB)
    inp = torch.from_numpy(low_img/127.5-1).permute(2,0,1).unsqueeze(0).float().to(device)
    out = (G(inp)[0].permute(1,2,0).cpu().numpy()+1)/2*255
    out = np.clip(out, 0, 255).astype(np.uint8)
# 直方图均衡化对比
he = cv2.cvtColor(cv2.imread('low.jpg'), cv2.COLOR_BGR2YUV)
he[:,:,0] = cv2.equalizeHist(he[:,:,0])
he = cv2.cvtColor(he, cv2.COLOR_YUV2BGR)
# 保存结果
cv2.imwrite('gan_result.jpg', cv2.cvtColor(out, cv2.COLOR_RGB2BGR))
cv2.imwrite('he_result.jpg', he)
print('已保存 gan_result.jpg 与 he_result.jpg')

运行与验证

准备一对256×256的成对图像，分别命名为low.jpg（低光照）和high.jpg（正常光照），置于脚本同目录。
直接运行脚本。训练约200轮（单张图像在GPU上约1分钟，CPU约5分钟）。
训练完成后，检查输出：gan_result.jpg为GAN增强结果，he_result.jpg为HE增强结果。

通过条件：GAN结果相比HE结果在暗部区域（如阴影、黑色物体）能恢复更多可见纹理与自然色彩，且无明显伪影或色偏。HE结果通常对比度过高、色彩失真，GAN结果应更接近真实光照下的观感。若GAN输出偏灰或偏暗，可增大感知损失权重（当前为10）或增加训练轮数。

常见问题：若训练不收敛（损失震荡剧烈），降低学习率至1e-4；若生成图像有棋盘伪影，检查反卷积层是否使用正确的kernel size与stride（当前配置已避免重叠）；若显存不足，将图像尺寸从256降至128并同步调整判别器输出尺寸。