基于CG-GAN的单张图像去雾实战：从原理到Python实现

1次阅读

共计 2338 个字符，预计需要花费 6 分钟才能阅读完成。

传统去雾算法如暗通道先验 (DCP) 在边缘设备上存在明显瓶颈：

计算复杂度高：DCP 需要求解软抠图（soft matting），时间复杂度达 $O(N^3)$，在移动端处理 1080P 图像需 500ms 以上
先验失效风险：天空区域或白色物体场景违反暗通道假设，导致色偏和光晕伪影
参数敏感性：大气光估计和透射率调整需要人工调参，难以适应多变天气条件

方法	PSNR(dB)	SSIM	推理速度(1080P)	参数量(M)
DCP	18.2	0.75	520ms	–
DehazeNet	21.7	0.83	120ms	0.8
CycleGAN	23.1	0.86	90ms	54.3
CG-GAN	24.5	0.89	35ms	23.7

测试环境：Intel i7-11800H + RTX 3060 Laptop

class ContextGuide(nn.Module):
    def __init__(self, in_ch=3):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, 64, 5, padding=2)
        self.down1 = Downsample(64, 128)  # 包含 MaxPool
        self.down2 = Downsample(128, 256)
        self.attn = CBAM(256)  # 通道 - 空间注意力
        self.up1 = Upsample(256, 128)
        self.up2 = Upsample(128, 64)
        self.out = nn.Conv2d(64, 3, 3, padding=1)

    def forward(self, x):
        x1 = F.relu(self.conv1(x))
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x3 = self.attn(x3)  # 关键上下文引导
        x = self.up1(x3) + x2  # 跳跃连接
        x = self.up2(x) + x1
        return torch.sigmoid(self.out(x))

采用 PatchGAN 结构，不同尺度处理：

原始分辨率：捕捉高频细节
1/ 2 下采样：平衡感受野
1/ 4 下采样：获取全局一致性

$$\mathcal{L}{total} = \lambda_1\mathcal{L}} + \lambda_2\mathcal{L{adv} + \lambda_3\mathcal{L}$$

感知损失：VGG16 relu3_3 特征图 MSE
对抗损失：Wasserstein GAN with Gradient Penalty
TV 正则项：抑制伪影 $\sum|\nabla x|^2$

def train_epoch(loader):
    G.train(); D.train()
    for hazy, clean in loader:
        # 数据增强
        hazy = augment(hazy)  # 随机翻转 + 颜色抖动

        # 生成去雾图像
        fake = G(hazy)

        # 判别器更新
        D.zero_grad()
        real_loss = -D(clean).mean()
        fake_loss = D(fake.detach()).mean()
        gp = gradient_penalty(D, clean, fake)
        loss_D = real_loss + fake_loss + 10*gp
        loss_D.backward()
        opt_D.step()

        # 生成器更新
        if step % 2 == 0:
            G.zero_grad()
            l_adv = -D(fake).mean()
            l_per = F.mse_loss(vgg(fake), vgg(clean))
            l_tv = TVLoss(fake)
            loss_G = 0.1*l_adv + l_per + 0.05*l_tv
            loss_G.backward()
            opt_G.step()

# 转换 ONNX
torch.onnx.export(G, hazy, "cggan.onnx", 
                 input_names=["input"],
                 output_names=["output"])

# 构建 TensorRT 引擎
builder = trt.Builder(logger)
network = builder.create_network()
parser = trt.OnnxParser(network, logger)
with open("cggan.onnx", "rb") as f:
    parser.parse(f.read())
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)  # 开启半精度
engine = builder.build_engine(network, config)