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

1次阅读

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

传统图像去雾方法如暗通道先验 (DCP) 存在两个致命缺陷：

计算复杂度高：DCP 需要求解软抠图方程，时间复杂度达 $O(N^2)$，在树莓派等边缘设备上处理 1080p 图像需要超过 2 秒
先验失效问题：当场景物体与大气光相似（如白色墙壁）或存在强光源时，暗通道假设不再成立，导致色彩失真

实测发现，DCP 在 RESIDE-SOTS 数据集上平均 PSNR 仅 18.7dB，且在移动端推理速度不足 5FPS

方法	PSNR(dB)	SSIM	推理速度(1080p)	参数量
DCP	18.7	0.82	0.5FPS	–
DehazeNet	21.3	0.86	15FPS	0.8M
CycleGAN	23.1	0.88	10FPS	11.4M
CG-GAN	24.6	0.91	25FPS	4.2M

（测试环境：RTX 3060 + PyTorch 1.10）

关键组件采用 UNet++ 结构，创新点在于：

上下文引导模块：在跳跃连接处加入轻量级 Transformer，计算流程为：

class ContextGuide(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.qkv = nn.Conv2d(channels, channels*3, 1)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        B, C, H, W = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=1)  # 拆分为 QKV
        q = q.view(B, C, -1).permute(0, 2, 1)  # [B, HW, C]
        attn = self.softmax(torch.bmm(q, k))   # 注意力权重
        out = torch.bmm(attn, v)               # 上下文聚合
        return out.view(B, C, H, W)

多尺度特征融合：在解码器每个上采样阶段引入 3×3 和 5×5 并行卷积核，通过 concatenation 融合不同感受野特征

总损失包含四项：

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

感知损失：使用 VGG16 的 relu3_3 层特征
对抗损失：采用 WGAN-GP 的梯度惩罚策略
TV 正则化：抑制输出图像的噪声

# 数据增强示例
train_transform = transforms.Compose([transforms.RandomCrop(256),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.1, 0.1, 0.1),
    transforms.ToTensor()])

# 多尺度判别器实现
class MultiscaleDiscriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.downsample = nn.AvgPool2d(3, stride=2, padding=1)
        self.models = nn.ModuleList([NLayerDiscriminator(3) for _ in range(3)
        ])

    def forward(self, x):
        outputs = []
        for model in self.models:
            outputs.append(model(x))
            x = self.downsample(x)
        return torch.cat(outputs, 1)

使用 TensorRT 进行 INT8 量化：

# 转换 ONNX
torch.onnx.export(model, dummy_input, "cggan.onnx", 
                 opset_version=11, 
                 do_constant_folding=True)

# TensorRT 优化
trt_cmd = f"trtexec --onnx=cggan.onnx --int8 --saveEngine=cggan.trt"
os.system(trt_cmd)

实测量化后模型在 Jetson Xavier 上推理速度从 15FPS 提升到 28FPS，PSNR 仅下降 0.3dB

训练振荡问题：
采用两时间尺度更新规则(TTUR)，生成器和判别器使用不同学习率
添加谱归一化 (Spectral Norm) 稳定训练
小样本训练：
使用 Foggy Cityscapes 等合成数据预训练

采用 CutMix 数据增强：

def cutmix(batch):
    lam = np.random.beta(1.0, 1.0)
    rand_index = torch.randperm(batch.size(0))
    target_a = batch
    target_b = batch[rand_index]
    h_start, w_start = random_pos(batch.size(2), batch.size(3))
    batch[:, :, h_start:, w_start:] = target_b[:, :, h_start:, w_start:]
    return batch, lam