深度学习项目-图片自动着色

简介

这是利用CNN分类模型结构进行图片自动着色的一个深度学习项目。本项目的实现参照Richard Zhang等人的论文Colorful Image Colorization实现。本项目代码包括模型均已开源于Github，部署教程见文末Github地址。

项目流程

整体思路

对一个Gray的图片建立到RGB的映射模型是一个比较复杂的回归问题，且效果难以把控。所以这里提出使用Lab颜色空间作为整个模型的映射机制，由于Lab将色彩信息放置在ab两个通道的特性，而L通道仅保存亮度等信息，这些信息可以从Gray图直接得到。

这就是说模型的输入是Lab中L通道的矩阵，目标是一个ab两通道的二层矩阵，建立的是从L向ab的映射关系。

当得到ab通道的信息再组合上L通道信息即得到彩色的Lab图片，而Lab图片到RGB图片的转换是固定的，也就是说，得到了期待的彩色RGB图片。

前期优化

那么，对训练过程中的一张图片，得到它的L信息很容易，得到ab信息也很容易，但是Lab色彩空间有个很大的问题就是其中包含了很多人类不可见的色彩，这些色彩是干扰模型预测的要素。在此，提出来一种新的思路，选取数据集中常见的色彩，组成ab值对，作为ab通道的信息。

实验证明，这种思路很成功，最后一共得到了313个常用值对信息。这就是说，预测的目标变为了一个值对或者313个值对的概率分布，原来复杂的回归变成了一个简单不少的分类问题。

当然，这里有一个问题，那就是训练时对于每个RGB图片（转为Lab图片），如何得到其目标的ab值对？这里使用了5近邻搜索法得到像素点的ab值对。

模型设计

既然是图片分类问题，如何保证分类的准确率呢？事实证明，VGGNet在图片分类上有着不错的效果，再加以近几年常用的BN层保证模型的训练强度，同时利用扩张卷积在两个block获得更大的感受野，就形成了下图的模型。

后续处理

得到了313个ab值对的颜色概率，如何将其映射回ab通道的色彩呢，在尝试了平均法和分立法之后（前者上色过于融合后者过于分界），使用兼有两种要求的模拟退火搜索法得到ab两个通道的矩阵。

损失函数

事实上，优化函数的研究这些年已经有了很大的发展，Adam已经可以满足本项目的需求了，但是在尝试使用分类经典的交叉熵函数时发现训练的模型上色均比较暗淡，这是因为imagenet数据集的图片普遍比较背景暗淡，为了产生比较鲜艳的色彩，提出了加上颜色再平衡系数的交叉熵函数，训练的模型效果得以改善。

模型构建

本项目模型构建基于Keras Function API（TensorFlow后端），训练与Tesla T4 16G GPU。模型提供于Github，可以使用Netron可视化h5模型。修改的交叉熵函数可见于脚本目录下的loss_function.py文件。

上色效果

这里提及一下，这里上色的目的从来不是多么接近原来的彩色图片（事实上需要应用模型的地方也不可能有对照的彩色图片，不然就没有上色的必要了），而是上色的结果能够让人眼看上去不觉得有什么异常即可。

上述所说的特点在图1体现明显。

补充说明

本项目代码和模型开源于我的Github，欢迎star或者fork。其中先验参数参考foamliu的项目。实现过程对论文模型最后两个block进行了稍加修改。博客同步至个人博客网站，欢迎查看。