灰度/红外图像彩色化

发表于 | 分类于 | 阅读次数：

【摘 要】 本项目实现两个任务：1、使用ImageNet数据集训练CartoonGan和Pix2Pix网络，使其能够将灰度图转化为彩色图，分析两者性能的差距。2、使用非成对红外图像数据集训练CartoonGan，使其能够将红外图像转化为灰度图像。

Pix2Pix论文链接：https://arxiv.org/abs/1611.07004

Pix2Pix的GitHub地址：https://github.com/phillipi/pix2pix

CartoonGan论文链接：点击跳转

CartoonGan的GitHub地址：https://github.com/znxlwm/pytorch-CartoonGAN

一、介绍

1.1 功能介绍

实现功能如下：

  • 使用ImageNet数据集训练CartoonGan和Pix2Pix网络，使其能够将灰度图转化为彩色图。
    • 对比CartoonGan和Pix2Pix的性能差距
  • 使用非成对红外图像数据集训练CartoonGan，使其能够将红外图像转化为灰度图像。

1.2 网络介绍

1.2.1 Pix2Pix

pix2pix1使用的是成对数据集进行训练，生成器采用的是Unet2结构，判别器使用的是马尔科夫性的判别器PatchGAN3。结构如下图所示：

1590650769339

结构比较简单，修改网络的数据集，既可以实现灰度图像彩色化。

细节可查看某大佬笔记

1.2.2 CartoonGan

CartoonGan4使用的是非成对数据集进行训练，生成器使用的编解码结构，判别器就是一个普通的分类网络，结构如下所示：

1590652379962

输入为图片为3通道，对数据集进行简单的预处理既可以开始训练。

细节可看CSDN大佬

二、实验步骤

  • 将下载好的ImageNet数据集灰度化。

2.1 Pix2Pix灰度图彩色化

  • 配置环境
  • 分割数据集，训练集0.7，测试集0.3。
  • 设置好目录后开始训练，在GTX 2070s上面训练3天，共200个epoch。

2.2 CartoonGan灰度图彩色化

  • 配置环境
  • 分割数据集，训练集0.7，测试集0.3。
  • 由于作者没有提供灰度图彩色化代码，需要自己关注输出结果调整网络参数。
    • 现象：图片无明显彩色化，怀疑是内容损失惩罚系数过大，降低内容损失为0.2

2.3 CartoonGan红外图彩色化

更换数据集为红外图像，红外图像和灰度图的区别在于红外图像更加模型，需要模型具备更强的想象力。步骤如下：

  • 更换数据集，其他和2.2一样。

三、实验结果

3.1 Pix2Pix灰度图彩色化

Pix2Pix输出结果如下图所示：左边为灰度图，中间为生成图，右边为对应彩色图。

1590656650202

3.2 CartoonGan灰度图彩色化

CartoonGan输出结果如下图所示：左边为灰度图，中间为生成图，右边为对应彩色图。

1590656777456

3.3 CartoonGan红外图彩色化

30_epoch_project_name_test_5

27_epoch_project_name_test_4

四、总结

除此之外，对比了峰值信噪比，得出如果有成对数据集，Pix2Pix性能优于CartoonGan。但是面对红外图像，无成对数据集的情况，Pix2Pix便无能为力。

在CartoonGan的训练过程中，作者默认设置的内容损失有些过大，调小会得到更好的结果。同时CartoonGan具有很大的成长空间，修改特征提取方式（作者使用VGG提取特征，权重为VGG官方默认）应该会得到一些有趣的结果。

Pix2Pix处理每张图片用时：0.04914107492991856 s

CartoonGan处理每张图片用时：0.02963963421908292 s

五、环境配置说明

只记录Pix2Pix，CartoonGan配置方式类似。电脑需要安装好Ubuntu18.04，cudnn等必备库

5.1 环境配置

  • python：3.6.6
  • pytorch：1.5.0
  • pillow：6.0.0
  • scipy：1.2.0
  • torchvision：2.3.1
  • visdom
  • dominate：2.3.1

将上面的环境先安装好，conda创建环境时命名为：pix2pix，后面提示缺少什么模块安装什么模块。

5.2 数据集配置

数据集下载，由于数据集太大，使用测试集作训练，总共40670张。

官网下载链接

百度云下载链接: ch9u

5.2.1 目录结构
  • 根目录
    • train
      • 放置30000张图片
    • test
      • 放置10670张图片

使用时，只需要将根目录地址给程序。

5.2.2 数据集分割

数据集分割代码，需要提前创建train和test文件夹，代码会自动将test2017下的全部图片，分割到train和test文件夹中。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import os
import random
import numpy as np
import shutil

path = '/media/w/文档/ubuntu/coco/test2017'
filelist = os.listdir(path)

random.shuffle(filelist)

for index, pathName in enumerate(filelist, 1):
    src = os.path.join(path, pathName)

    if index <= 30000:
        dst = os.path.join('train/', pathName)
    else:
        dst = os.path.join('test/', pathName)

    shutil.copyfile(src, dst)

至此，数据集已处理完毕

5.3 打开visdom服务

必须要打开，不然测试或运行的时候会报错！！！！！

打开终端，进入程序根目录下

  • 激活环境
1
conda activate pix2pix
  • 打开visdom服务
1
python -m visdom.server

开始自动下载脚本，之后会弹出链接，点击链接即可在浏览器查看训练情况

执行后，出现以下提示表示正在下载，如果太久，参考该链接解决: https://github.com/casuallyName/document-sharing/tree/master/static

1
2
Checking for scripts.
Downloading scripts, this may take a little while

提示：visdom在ubuntu下的地址为：

1
~/miniconda3/envs/pix2pix/lib/python3.6/site-packages/visdom

5.4 训练配置

  • 下载代码，进入代码根目录
1
git clone https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix

  • 安装环境

    • python 3.6.6
    • pytorch 1.5.0
    • visdom
    • 其他的提示缺啥，安装啥。

  • 输入以下命令开始训练

1
python train.py --dataroot /media/w/文档/ubuntu/coco --name BW2COLOR --model colorization

其中dataroot是数据集地址，数据集参考第二节的做法。

5.5 测试说明

在代码根目录下创建weight/BW2COLOR文件夹，将权重文件： 0awi放入

1
python test.py --dataroot /media/w/文档/ubuntu/dataset/imagenet/ --name BW2COLOR --model colorization --checkpoints_dir weight

其中dataroot是测试集地址，需要把测试图片放入到imagenet内的test文件夹中。

5.6 信噪比比较

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import numpy as np
import math
from PIL import Image

def PSNR(img1, img2):
    mse = np.mean((img1 - img2)** 2)
    if mse == 0:
        return 100
    PIXEL_MAX = 255.0
    return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))

img1 = Image.open("./results/BW2COLOR/test_latest/images/000000000183_fake_B_rgb.png")
img2 = Image.open("./results/BW2COLOR/test_latest/images/000000000183_real_A.png")

img1_array = np.array(img1)
img2_array = np.array(img2)

print("The PSNR between image1 and image2 is: %.3f" % PSNR(img1_array,img2_array))

参考文献

1. Image-to-Image Translation with Conditional Adversarial Networks ↩
2. U-Net: Convolutional Networks for Biomedical Image Segmentation ↩
3. Patch-GAN Image-to-Image Translation with Conditional Adversarial Networks ↩
4. CartoonGAN: Generative Adversarial Networks for Photo Cartoonization ↩
-------------本文结束感谢您的阅读-------------
0%