老旧黑白片修复机——使用卷积神经网络图像自动着色实战（原文附PyTorch代码）

人工智能和深度学习技术逐渐在各行各业中发挥着作用，尤其是在计算机视觉领域，深度学习就像继承了某些上帝的功能，无所不能，令人叹为观止。照片承载了很多人在某个时刻的记忆，尤其是一些老旧的黑白照片，尘封于脑海之中，随着时间的流逝，记忆中对当时颜色的印象也会慢慢消散，这确实有些可惜。但随着科技的发展，这些已不再是比较难的问题。在这篇文章中，将带领大家领略一番深度学习的强大能力——将灰度图像转换为彩色图像。文章使用PyTorch从头开始构建一个机器学习模型，自动将灰度图像转换为彩色图像，并且给出了相应代码及图像效果图。整篇文章都是通过iPython Notebook中实现，对性能的要求不高，读者们可以自行动手实践一下在各自的计算机上运行下，亲身体验下深度学习神奇的效果吧。 PS：不仅能够对旧图像进行着色，还可以对视频（每次对视频进行一帧处理）进行着色哦！闲话少叙，下面直接进入正题吧。

简介

在图像着色任务中，我们的目标是在给定灰度输入图像的情况下生成彩色图像。这个问题是具有一定的挑战性，因为它是多模式的——单个灰度图像可能对应许多合理的彩色图像。因此，传统模型通常依赖于重要的用户输入以及输入的灰度图像内容。 最近，深层神经网络在自动图像着色方面取得了显着的成功——从灰度到彩色，无需额外的人工输入。这种成功的部分原因在于深层神经网络能够捕捉和使用语义信息（即图像的实际内容），尽管目前还不能够确定这些类型的模型表现如此出色的原因，因为深度学习类似于黑匣子，暂时无法弄清算法是如何自动学习，后续会朝着可解释性研究方向发展。 在解释模型之前，首先以更精确地方式阐述我们所面临的问题。

问题

我们的目的是要从灰度图像中推断出每个像素（亮度、饱和度和色调）具有3个值的全色图像，对于灰度图而言，每个像素仅具有1个值（仅亮度）。为简单起见，我们只能处理大小为256 x 256的图像，所以我们的输入图像大小为256 x 256 x 1（亮度通道），输出的图像大小为256 x 256 x 2（另两个通道）。 正如人们通常所做的那样，我们不是用RGB格式的图像进行处理，而是使用LAB色彩空间（亮度，A和B）。该色彩空间包含与RGB完全相同的信息，但它将使我们能够更容易地将亮度通道与其他两个（我们称之为A和B）分开。在稍后会构造一个辅助函数来完成这个转换过程。 此外将尝试直接预测输入图像的颜色值（即回归）。还有其他更有趣的分类方法，但目前坚持使用回归方法，因为它很简单且效果很好。

数据

着色数据无处不在，这是由于我们可以从任何一张彩色图像中提取出灰度通道。对于本文项目，我们将使用MIT地点数据集中的一个子集，该子数据集包含地点、景观和建筑物。

工具

本文使用PyTorch构建和训练搭建的模型。此外，我们还了使用torchvision工具，该工具在PyTorch中处理图像和视频时很有用，以及使用了scikit-learn工具，用于在RGB和LAB颜色空间之间进行转换。

模型

模型采用卷积神经网络构建而成，与传统的卷积神经网络模型类似，首先应用一些卷积层从图像中提取特征，然后将反卷积层应用于高级（增加空间分辨率）特征。 具体来说，模型采用的是迁移学习的方法，基础是ResNet-18模型，ResNet-18网络具有18层结构以及剩余连接的图像分类网络层。我们修改了该网络的第一层，以便它接受灰度输入而不是彩色输入，并且切断了第六层后面的网络结构：

训练

损失函数

由于使用的是回归方法，所以使用的仍然是均方误差损失函数：尝试最小化预测的颜色值与真实（实际值）颜色值之间的平方距离。

由于问题的多形式性，上述损失函数对于着色有一点小的问题。例如，如果一件灰色的衣服可能是红色或蓝色，而模型若选择错误的颜色时，则会受到严厉的惩罚。因此，构建的模型通常会选择与饱和度鲜艳的颜色相比不太可能“非常错误”的不饱和颜色。关于这个问题已经有了重要的研究（参见Zhang等人），但是本文将坚持这种损失函数，就是这么任性。

优化

使用Adam优化器优化选定的损失函数（标准）。

加载数据

使用torchtext来加载数据，由于我们需要LAB空间中的图像，所以首先必须定义一个自定义数据加载器（dataloader）来转换图像。接下来，对训练数据和验证数据定义变换。

辅助函数

在进行训练之前，定义了辅助函数来跟踪训练损失并将图像转换回RGB图像。

验证

在验证过程中，使用torch.no_grad（）函数简单地运行下没有反向传播的模型。

训练

在训练过程中，使用loss.backward（）运行模型并进行反向传播过程。

预训练模型

如果你想运用预训练模型而不想从头开始训练的话，我已经为你训练了好一个模型。该模型在少量时间内接受相对少量的数据训练，并且能够工作正常。可以从下面的链接下载并使用它

结果

有趣的内容到了，让我们看看深度学习技术实现的效果吧！

结论

在这篇文章中，使用PyTorch工具从头创建了一个简单的自动图像着色器，没有太复杂的代码，只需要简单的准备好数据并设计好合理的模型即可得到令人令人兴奋的结果，此外，这仅仅只是起步，后续还有很多地方可以进行改进优化并进行推广。 如果你对这类任务还感兴趣的话，可以继续阅读以下资源：

• Incorporating user feedback into colorization: Zhang et al., 2017
• Using colorization for unsupervised learning: Larsson et al. , 2017
• Add it to a camera app, like Google just did
• Colorize a scene from an old film, as I experimented with 本文构建的模型完整代码，或者关于彩色化的更详细的技术报告已经上传至GitHub，读者可以在该网站上查看整个完整的项目文档。赶紧动手尝试一下吧，感兴趣的可以继续进行优化或完善，开发出一款APP等，说不定下一个风投就是你哦！