AOD-Net预训练模型实战：从图像去雾到工业部署的完整解决方案

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

背景痛点

传统图像去雾方法如暗通道先验（DCP）在理想光照条件下表现尚可，但在实际工业场景中面临严峻挑战：

1. 复杂光照适应性差：工厂雾霾常伴随强点光源或低照度，导致透射率估计失效
2. 实时性瓶颈：DCP 的软抠图算法单帧处理需 200ms+，无法满足产线检测的 50ms 延迟要求
3. 参数固化问题：大气散射模型中的 β、A 等参数需要人工标定，难以应对动态雾浓度变化

模型解析

AOD-Net 通过端到端学习实现大气光与透射率的联合优化，其创新架构如下图所示：

核心设计亮点：

1. 轻量级 K 估计模块：
2. 采用 5 层卷积实现透射率图预测（仅 0.8M 参数）
3. 最后一层使用 Sigmoid 约束输出到 (0,1) 区间

4. 数学表达：K(x)=1-t(x)，其中 t(x)为透射率

5. 大气光联合优化：

6. 将全局大气光 A 作为可学习参数
7. 通过全连接层实现 A∈R^3 的预测
8. 去雾公式：J(x)=[I(x)-A]/K(x)+A

参数量对比（输入 512×512）：

代码实战

数据加载器实现

class FoggyCityscapes(Dataset):
    def __init__(self, root, transform=None):
        self.image_pairs = [...]  # 加载图像对路径
        self.transform = transform

    def __getitem__(self, idx):
        clear_img = cv2.imread(self.clear_paths[idx])
        foggy_img = cv2.imread(self.foggy_paths[idx])

        # 自适应直方图均衡化
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        foggy_img = cv2.cvtColor(foggy_img, cv2.COLOR_BGR2LAB)
        foggy_img[...,0] = clahe.apply(foggy_img[...,0])
        foggy_img = cv2.cvtColor(foggy_img, cv2.COLOR_LAB2BGR)

        if self.transform:
            clear_img = self.transform(clear_img)
            foggy_img = self.transform(foggy_img)

        return foggy_img, clear_img

混合损失函数

def hybrid_loss(pred, target):
    # L1 损失
    l1_loss = F.l1_loss(pred, target)

    # SSIM 损失（窗口尺寸 11×11）ssim_loss = 1 - ssim(pred, target, window_size=11, size_average=True)

    # 梯度裁剪
    total_loss = l1_loss + 0.5*ssim_loss
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    return total_loss

TensorRT INT8 量化

# 转换 ONNX
torch.onnx.export(model, dummy_input, "aodnet.onnx")

# 构建 TensorRT 引擎
builder = trt.Builder(logger)
network = builder.create_network()
parser = trt.OnnxParser(network, logger)

# INT8 校准
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.INT8)
calibrator = Int8EntropyCalibrator(calib_data)
config.int8_calibrator = calibrator

engine = builder.build_engine(network, config)

性能优化

在 NVIDIA T4 显卡上的测试结果：

避坑指南

1. 雾浓度估计不准：
2. 在自定义数据上微调 K 估计模块

3. 添加雾密度分类头辅助训练

4. 多相机批处理：

5. 使用 NVIDIA DALI 加速数据加载

6. 对不同分辨率输入进行 zero-padding 到统一尺寸

7. 动态分辨率处理：

8. 预分配最大尺寸的显存池
9. 采用 TRT 的动态 shape 特性

   profile = builder.create_optimization_profile()
   profile.set_shape("input", (1,3,480,640), (1,3,720,1280), (1,3,1080,1920))

开放性问题

如何将 AOD-Net 适配到水下图像增强场景？考虑以下方向：
– 修改物理模型：将大气散射模型替换为水下光传播模型
– 数据增强策略：模拟不同水深的光吸收特性
– 多任务学习：联合估计色偏校正与去雾