基于深度学习cGAN模型进行图像去雾

基于深度学习cGAN模型进行图像去雾摘要相关工作提出模型实验成果代码运行环境要求使用步骤（按照实际遇到的问题进行了修改，与官方给出的有点不同）官方Replicating步骤

参考论文：N. Bharath Raj and N. Venketeswaran, "Single Image Haze Removal using a Generative Adversarial Network," 2020 International Conference on Wireless Communications Signal Processing and Networking (WiSPNET), Chennai, India, 2020, pp. 37-42, doi: 10.1109/WiSPNET48689.2020.9198400.

摘要

本方法选自WiSPNET2020的一篇关于使用深度学习进行图像去雾的文章。传统的从图像中去除雾霾的方法依赖于对大气透射等参数的估计，而当处理单一图像时，由于缺乏深度信息，这成为一个不适定的问题。在本文中，作者提出了一种基于端到端学习的方法，利用改进的条件生成对抗网络（Conditional GAN, cGAN）直接去除图像中的雾霾。文中采用Tiramisu模型代替经典的U-Net模型作为生成器，因为它具有更高的参数效率和性能；同时还采用了基于区域的判别器来减少输出中的伪像。此外，为了进一步提高输出的感知质量，作者设计了一种混合加权损失函数来训练模型。通过采集真实图像进行训练、验证、测试以及与传统方法对比，该模型通过深度学习方法显现出了一定的优势。

本文主要有以下四点成果：

介绍使用cGAN直接去除给定单个场景图像的雾霾的方法。
用56层Tiramisu模型代替U-Net架构。后者是参数有效的，并且在语义分割方面具有最先进的性能。
在cGAN中使用块判别器以减少伪像。
加权损失函数除了标准的cGAN生成器损失外，还考虑了L1损失和感知损失的影响，以激励模型产生视觉上贴合人感知的输出图像。

提出模型

图 2 提出模型的生成器结构

上图为cGAN模型生成器的主要结构，各模块的联合作用与U-Net相似。如图所示，它由密度块（DB），向下过渡层（TD），向上过渡层（TU）和瓶颈层组成。另外，它包含一个输入卷积和一个输出卷积。虚线表示的是串联操作。

图 3 Dense Block (DB) 的结构图

其中，该模型在编码侧有5个密度块，在解码侧也有5个密度块，并且有一个密度块作为瓶颈层。如图是4层密度块的示意图。第一层用来从输入图像中提取出k个特征图，并将这些特征图连接到输入图像中。然后第二层用来提取出另外的k个特征图，这些特征图又被连接到之前的特征图中。重复以上操作4次之后，密度块的输出是4层输出的串联，因此包含4 * k个特征图。

图 4 生成器各复合层结构

在编码侧，每个密度块（DB）后面都有一个向下过渡（TD）层。TD层包括批处理标准化，激活函数，卷积，丢弃和平均池化操作。类似地，在解码侧，每个密度块（DB）后面都有一个向上过渡（TU）层。 TU层仅具有转置卷积操作。编码侧和解码侧中的每个DB层都有4个复合层。瓶颈层中的DB层包含15个复合层。每个复合层都包含批处理标准化，激活函数，卷积和丢弃操作。从空间上说，穿过TD层后，图像的空间尺寸减半，而穿过TU层后，图像的空间尺寸加倍。

以下是各层在生成器中发挥的作用：

BN层：加快网络的训练和收敛的速度，控制梯度爆炸防止梯度消失，防止过拟合。
ReLU层：增加了神经网络各层之间的非线性关系，否则，如果没有激活函数，层与层之间是简单的线性关系，每层都相当于矩阵相乘。
Dropout层：防止过拟合。
Average Pooling层：下采样，降维、去除冗余信息、对特征进行压缩、简化网络复杂度、减少计算量、减少内存消耗等等，实现非线性，可以扩大感知野，可以实现不变性。
卷积层：下采样。
转置卷积层：上采样，从低维度进入高维度。

图 5 提出模型的判别器结构

上图为cGAN模型判别器的主要结构。它的输入包含一张带雾的图像和该图像对应的一张生成网络生成的有效图像。它输出一个30x30的矩阵，用于测试图像的真实性。同样的，各层的作用和生成器中所对应的层级是一样的。

图 6 不同尺寸区域判别器下的判别效果

图 7 不同尺寸区域判别器下的判别分数

该网络对目标图像和生成的图像执行逐块比较，而不是逐像素进行比较。将图像放入CNN来实现这一过程，该CNN输出的接收区域大于一个像素（即对应于原始图像中的像素斑块）。我们使用的是70x70区域判别器，并且在最后一组特征图上使用逐像素比较。特征图上的有效接收区域大于一个像素，因此它覆盖了图像的一部分，这消除了图像中的大量伪像。

1x1 PixelGAN鼓励更大的颜色多样性，但对空间统计没有影响。16x16 PatchGAN可以产生局部清晰的结果，但也会导致超出其可观察范围的平铺伪像。70x70 PatchGAN会在空间和色彩维度上使输出即使是不正确的也要保持清晰。完整的286x286 ImageGAN产生的结果在视觉上与70x70 PatchGAN相似，但根据FCN网络评分标准来判断其质量略低。

相比于直接使用标准生成器损失函数，我们更倾向用L1损失和感知损失函数对其进行增强，最终组成一个加权生成器损失函数LT。这样做是为了提高生成图像的质量。我们直接将LT最小化作为目标来训练我们的生成器，公式如下。

其中，给定一个模糊图像集X及其对应的无雾图像集Y，标准生成器目标是使LG最小化，公式如下。

另外，使用加权的L1损失可减少输出图像中的伪像。目标图像y与生成的图像（）之间的L1损耗计算如下。

同时，我们将特征重建感知损失添加到了总损失中。但是，我们不是使用L2，而是使用按常数C缩放的均方误差（MSE）。生成的图像和目标图像将通过不可训练的VGG-19网络传递，然后计算两个图像输出的MSE损失。

而判别器的目标是最大化以下目标损失函数LD。

实验成果

以下是相关数据及参数的设定情况：

Dataset: NYU Depth V2 dataset (with the size [256, 256, 3]).
Training Details:
Performance Metrics:

图 8 带雾图像与无雾图像比较示例

除雾算法的性能可以从几个因素进行评估，其中两个最常用的因素是PSNR和SSIM。峰值信噪比（PSNR）衡量算法从噪声图像中去除噪声的能力。两个相同的图像将具有无穷大的PSNR值。

为了测量除雾能力，较高的PSNR值表示较好的性能。结构相似性指数度量（SSIM），用于度量两个图像的相似度。两个相同的图像的SSIM为1。

相对而言，PSNR高的图像在视觉上无法保证令人感知质量良好，而SSIM高的图像能够代表更好的去雾质量。因此，我们需要一个同时具有这两个属性的指标。我们定义了一个称为Score的度量值，它是图像的PSNR和SSIM的加权和。

图 9 训练过程中网络的性能感知比较

图 10 训练过程中网络的性能数据比较

以上是训练模型的过程中根据当前阶段网络对图像进行处理，并与原图对比之后所得到的性能数据比较。可以看到，网络在训练过程中不断对图像进行特征提取，包括图中物体轮廓、颜色、亮度、对比度等因素，并与无雾图像（也就是目标图像）进行比较得出相应的分数。其中，G_D_Step代表的是当前经历的生成器+判别器所组成循环的次数。本文中抽取十次优化情况用作比较，可以从图像中直观地看出，该网络正在逐步的对图像的重要特征进行逐步提取，所生成的图像也越来越相似于目标图像。而从量化数据也可以看出，随着训练次数的迭代，生成图像与目标图像对比所得出的Score、PSNR、SSIM都在逐渐增加。

图 11 带雾未处理图像、经过epoch=5的神经网络图像、经过epoch=20的神经网络图像以及原图像对比

图 12 以上三类图像与原图对比后的Score、PSNR、SSIM值

以上是带雾未处理图像、经过epoch=5的神经网络图像、经过epoch=20的神经网络图像以及原图像对比的情况，以及以上三类图像与原图对比后的Score、PSNR、SSIM值。可以看到，相比于原图来说，经过训练的GAN模型处理后的图像无论是感官上还是量化数据上都有了质的飞跃，特别是Score、PSNR、SSIM三个值都接近翻倍，可见通过该模型对带雾图像进行去雾的效果较为明显。值得注意的是，以上也对经历了5个及20个epoch的模型进行了比较，可以看到后者处理得到的图像在细节上还是更为精确、轮廓及色彩较为明显的，而且其Score、PSNR、SSIM都有小幅的上涨。

比较可惜的是，由于带雾图像的确比较模糊，在深度较大的地方该模型还是不能准确的还原出图像中少部分物体的准确颜色及轮廓，这可能是由于深度较大的地方其雾污染已使环境中部分视野完全受限（具体表现为该像素区域显示接近全白），从而导致信息量缺失，后续可以考虑通过强化轮廓等基本特征的作用效果进行性能优化。

代码运行

参考代码：https://github.com/thatbrguy/Dehaze-GAN

环境要求

TensorFlow (version 1.4+)，代码是基于TensorFlow 1.x编写的，现在tensorflow.org发布的最新版本已经是2.x的，所以通过pip等方式安装一般都是2.x版本，且1.x的版本通道已经被关闭了，接下来我会讲解如何将1.x版本代码过渡到2.x版本。
Matplotlib
Numpy
Scikit-Image
NVIDIA GPU驱动 418.x以上 + CUDA Toolkit 10.1 + cuDNN SDK 7.6 （版本不能错）

由于本人用的是实验室台式电脑加4块1080Ti的小型服务器，不太满意跑的速度，最后是用pycharm连接租的云GPU跑的，硬件环境不太好的朋友可以参考。

使用步骤（按照实际遇到的问题进行了修改，与官方给出的有点不同）

下载github的源码，可以手动下载zip，也可以git：
```
git clone https://github.com/thatbrguy/Dehaze-GAN.git
```
需要下载一个VGG-19模型用于跑出感知损失，作者用的是Machrisaa的版本，点击此处下载。
下载数据集，我用的是和作者一样的NYU Depth Dataset V2数据集。可以点击此处手动下载数据集，然后把数据集放到项目文件夹中运行extract.py提取图像（注意修改代码中mat文件路径），也可以在项目文件夹按照作者指引执行命令：
```
wget -O data.mat http://horatio.cs.nyu.edu/mit/silberman/nyu_depth_v2/nyu_depth_v2_labeled.mat
python extract.py
```
最终得到的是项目文件夹中A和B两个文件夹，分别存放带雾图像和对应的无雾图像。当然这一步也可以用你自己的数据集，只需要按要求把图分好放在两个文件夹并确保图的尺寸是(256, 256, 3)。
（重点）如上文所说，我们需要将TensorFlow 1的代码迁移到TensorFlow 2适用的格式，这点使用TensorFlow 1的朋友可以跳过。在参考官方文档《Migrate your TensorFlow 1 code to TensorFlow 2》之后得到以下较为简单的方法（直接修改代码的方法可以参照文档）。
运行自动upgrade到2.x的命令行：
```
tf_upgrade_v2 --intree my_project/ --outtree my_project_v2/ --reportfile report.txt
```
其中my_project处填需要转换的项目文件夹路径，my_project_v2处填存放2.x版本项目的文件夹路径，report.txt会记录项目中代码被修改和需要注意的地方。
另外，这个命令也不是万能的，一些没有被迁移到2.x的函数和库就需要我们进行手动替换，这个项目中只有一个需要替换的地方，那就是operations.py中的BN层函数用了contrib库的函数。通过参考官方文档，我们可以用2.x现有库中的函数对不能使用的函数进行替换，过程如下：
原代码：
```
def BatchNorm(input_, isTrain, name='BN', decay = 0.99):
    with tf.compat.v1.variable_scope(name) as scope:
        return tf.contrib.layers.batch_norm(input_, is_training = isTrain, decay = decay)
```
修改后代码：
```
def BatchNorm(input_, isTrain, name='BN', decay = 0.99):
    with tf.compat.v1.variable_scope(name) as scope:
        # return tf.contrib.layers.batch_norm(input_, is_training = isTrain, decay = decay)
        return tf.compat.v1.layers.batch_normalization(inputs=input_, momentum=0.99, training=True)
```
如果还有其他问题的话可以根据报错提示进行修正。
运行以下命令行训练模型（参考main.py可以查看各参数说明，务必确保训练模型时--mode=train --A_dir=A --B_dir=B --custom_data=true，如果--custom_data=false则会因为没有预存好的npy文件而报错，其他更改可选）：
```
python main.py --A_dir=A --B_dir=B --batch_size=2 --epochs=20 --custom_data=true --mode=train --save_samples=false
```
一个比较有用的功能是在模型训练的过程中，我们可以在每次保存checkpoint的时候输出当前模型跑出来的图像情况，实现的方法是在项目文件夹中新建文件夹命名为samples（名字可变），里面存放符合格式的带雾图像文件，然后命令行修改一下以下参数：--save_samples=true --sample_image_dir=samples，然后执行。
在训练完模型后我们可以在文件夹model（或者你修改的model_name对应的文件夹）中的checkpoint找到训练结果，默认保存的是后三次的checkpoint。注意因为代码编写的问题，打开checkpoint文件会发现里面的路径可能会被保存为绝对路径，为了避免后续读取checkpoint失败可以将路径改为相对路径。
运行以下命令行测试模型（同样要参考main.py可以查看各参数说明）：
```
python main.py --A_dir=input_dir --B_dir=GT_dir --mode=test
```
其中input_dir为存放待处理带雾图像的文件夹，GT_dir为存放对应无雾图像的文件夹，该命令会将带雾图像放到模型中处理，然后把处理后图像与ground truth图像进行对比得到Score、PSNR、SSIM显示在控制台。
运行以下命令行应用模型（同样要参考main.py可以查看各参数说明）：
```
python main.py --A_dir=input_dir --B_dir=result_dir --mode=inference
```
其中input_dir为存放待处理带雾图像的文件夹，result_dir为存放对应的处理后图像的文件夹，该命令会将带雾图像放到模型中处理，然后把处理后图像存放到result_dir中。

为了计算本地一些用其他方法跑出来的图片对应的Score，我根据test函数修改了一版test_local函数，可以直接计算本地图片之间的Score、PSNR、SSIM值，将代码放到inference函数后面，然后新增main.py选项、运行命令行时修改参数--mode=test_local就行，代码如下：
test_local函数代码：

    def test_local(self, input_dir, Real_dir):

        total_ssim = 0
        total_psnr = 0
        psnr_weight = 1/20
        ssim_weight = 1

        Real_list = os.listdir(Real_dir)
        input_list = os.listdir(input_dir)

        for i, (img_file, Real_file) in enumerate(zip(input_list, Real_list), 1):

            img = cv2.imread(os.path.join(input_dir, img_file), 1).astype(np.uint8)
            Real = cv2.imread(os.path.join(Real_dir, Real_file), 1).astype(np.uint8)

            print('Test image', i, end = '\r')

            psnr = compare_psnr(Real, img)
            ssim = compare_ssim(Real, img, multichannel = True)

            total_psnr = total_psnr + psnr
            total_ssim = total_ssim + ssim

        average_psnr = total_psnr / len(Real_list)
        average_ssim = total_ssim / len(Real_list)

        score = average_psnr * psnr_weight + average_ssim * ssim_weight

        line = 'Score: %.6f, PSNR: %.6f, SSIM: %.6f' %(score, average_psnr, average_ssim)
        print(line)

main.py执行函数修改：

if __name__ == '__main__':

    args = parser.parse_args()
    net = GAN(args)
    if args.mode == 'train':
        net.train()
    if args.mode == 'test':
        net.test(args.A_dir, args.B_dir)
    if args.mode == 'inference':
        net.inference(args.A_dir, args.B_dir)
    # New part below for test_local function
    if args.mode == 'test_local':
        net.test_local(args.A_dir, args.B_dir)

官方Replicating步骤

为了让每位用户都能亲身体验到最后模型的处理结果，作者提供了已经训练好的checkpoint文件以及相应的数据集，只需要根据github中指引下载到legacy文件夹并解压使用就行，此处不详述。

最后想啰嗦一句的是在实际跑代码的过程中我本人也有遇到挺多小问题，但基本上只要不断通过看错误提示和debug都能够找到解决办法，大部分都是tensorflow版本不兼容、cuda运行不正常、路径或文件格式不正确等等。另外神经网络的可扩展性还是挺高的，可以多多尝试。