U-Net深度学习灰度图像的彩色化

本文介绍了使用深度学习训练神经网络从单通道灰度图像生成3通道彩色图像的实验。在我看来，虽然它们因主题而异，但结果令人震惊，模型产生了幻觉，原始主题应该是什么颜色。

将灰度1通道图像增强为3通道彩色图像。

Div2k（https://data.vision.ee.ethz.ch/cvl/DIV2K/）训练数据集用于训练，本文中的这些图像基于Div2K验证集。
这是基于Fastai深度学习课程中建议和教授的技术。
尝试着色灰度图像的灵感来自于来自Fastai学生/校友，Jason Antic和他的DeOldify模型的GAN评论家训练模型：https：//github.com/jantic/DeOldify
此方法使用以下内容：

具有类似于DenseNet的交叉连接的U-Net架构
基于ResNet的编码器和基于ResNet的解码器
像素随机扩展ICNR初始化
基于VGG-16模型激活，像素损失和克矩阵损失的损失函数
从预训练的ImageNet模型转移学习

该模型或数学函数具有超过4000万个参数或系数。

损失函数真正使得模型的预测非常好，允许模型中的可训练参数接近最佳值。这几乎与我的超分辨率实验中使用的损失函数相同。

通过使用在训练期间静态而不是作为生成性对抗网络（GAN）批评者改变的损失函数，可以更快地收敛到给定体系结构的近似最佳训练模型。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

建筑细节

基于U-Net的学习者和特征丢失函数的体系结构非常复杂，因为它背后的原因。由于这些是我在不同实验中使用的架构和损失函数，我将它们分成两篇文章：
基于U-Net的模型架构：https：
//towardsdatascience.com/u-nets-with-resnet-encoders-and-cross-connections-d8ba94125a2c
用于训练网络的功能丢失功能：https：//towardsdatascience.com/loss-functions-based-on-feature-activation-and-style-loss-2f0b72fd32a9

ResNet-34编码器

ResNet-34是34层ResNet架构，在U-Net的下采样部分（U的左半部分）用作编码器。

解码器

具有编码器架构的Fastai动态U-Net学习器将自动构建U-Net架构的解码器侧，在这种情况下将ResNet-34编码器转换为具有交叉连接的U-Net。

预训练编码器

为了让模型知道如何进行图像改进，它大大加快了使用预训练模型的训练时间，因此模型具有需要检测和改进的特征类型的起始知识。使用在ImageNet上预先训练过的模型和权重几乎是理想的。用于pyTorch的预训练ResNet-34可从Kaggle获得：https：//www.kaggle.com/pytorch/resnet34

培训细节

数据集

Div2k数据集提供800个训练图像和100个验证图像。我在训练过程中使用了一定比例的训练图像进行验证，使Div2K验证集与训练过程完全隔离，以避免模型在训练期间观察到这些错误。
每个图像缩小为一个通道，用作灰度输入，原始RGB图像用作目标/地面实况。
培训数据进一步增强：

采取随机作物
水平翻转图像
调整图像的光照
添加透视变形

训练过程从如上所述的模型开始：基于ResNet-34架构的U-Net使用基于ImageNet上预训练的VGG-16架构的损耗函数预先训练在ImageNet上，并结合像素丢失和克矩阵。

功能丧失功能

损耗函数使用来自VGG-16模型的激活，像素损失和克矩阵损失。
像素损失是对预测中每个像素颜色与目标/地面实况中的每个像素颜色的接近程度的简单像素比较。
VGG-16型号的激活允许损失功能确定特征是否正确，在这种情况下它们看起来是否正确。
克损失着眼于生成图像的艺术风格，其特征以及颜色方面与原始图像的预测接近程度。
这些组合允许损失函数检查上下文中的精细细节和特征样式，以改进模型的预测。
这允许模型学习令人信服的颜色以应用于构成所生成的预测图像的特征。
关于功能和样式丢失功能的文章中描述了这种损失函数的更多细节：https：//towardsdatascience.com/loss-functions-based-on-feature-activation-and-style-loss-2f0b72fd32a9

训练模型的头部和主干

这里使用的三种方法特别有助于培训过程。这些是逐步调整大小，冻结然后解冻主干中权重的梯度下降更新和判别学习率。
该模型的架构分为两部分，主干和头部。
骨干网是U-Net的左手部分，是基于ResNet-34的网络的编码/下采样部分。头部是U-Net的右手部分，是网络的解码器/上采样部分。
主干基于在ImageNet上训练的ResNet34进行了预训练，这是转移学习。
头部需要其权重训练，因为这些层的权重被随机初始化以产生期望的最终输出。在一开始，来自网络的输出基本上是像素Shuffle子卷积之外的像素的随机变化，其中ICNR初始化被用作网络的解码器/上采样路径中的每个高级中的第一步。
网络主干中的权重被冻结，以便最初仅训练头部中的权重。

渐进式调整大小

最初训练大量较小的图像然后扩大网络和训练图像的速度更快。在512 px x 512px图像上执行该操作并且在更大的数据集上执行操作要快得多，这对于64px到64px图像进行着色是一项更容易的任务。这称为渐进式调整大小，它还有助于模型更好地概括，因为可以看到更多不同的图像，并且不太可能过度拟合。
对于每个图像训练大小，训练模型直到训练损失和验证损失改进趋于平稳，然后图像数据尺寸加倍。动态U-Net学习器处理较大的输入和输出图像。

训练色彩64 x 64像素的图像

较大的批量大小加快了小图像的训练。

花了50个时间进行训练才能显着停止改善。由于模型需要更多像素/数据来预测特征的颜色和颜色样式，因此损失显着改善。

训练色彩64 x 64像素的图像

训练色彩为128 x 128像素的图像

在25个训练时期之后，损失显着改善。同样，由于模型需要更多像素/数据，损失停止显着改善。训练图像尺寸加倍。

训练色彩为128 x 128像素的图像

训练色彩256 x 256像素的图像

在10个训练时期之后，损失显着改善，训练大小增加到512×512像素。同样，由于模型需要更多像素/数据，损失停止显着改善。

训练色彩256 x 256像素的图像

训练色彩512 x 512像素的图像

在10个训练时期之后，由于GPU内存限制，损失停止显着改善，增加了训练图像大小并且训练转移到解冻整个网络的权重。

训练色彩512 x 512像素的图像

训练512x512像素图像，骨干重量解冻

现在受过训练，头部顶部的头部允许模型学习与骨干中的预训练知识做一些不同的事情。
然后解冻整个模型的权重，并用有辨别力的学习率训练模型。这些学习率在第一层组中小得多，然后在第二层组中增加并且在头部（最后一层）组中再次增加。
在10个训练时期之后，损失显着改善。
预测中的颜色看起来更接近目标/地面事实。

训练512 x 512像素图像，骨干重量解冻

结果

该结果来自在Div2K验证集上从训练模型进行预测。
如果颜色与目标和地面实况不匹配的结果，如果它们具有说服力，则不一定是错误的预测，如果目标/地面实况图像无法比较，则可能被认为是正确的。
算法无法知道该图像中的原始颜色是什么。该模型基本上可以将其训练的知识与基于图像中的内容，图像中的特征和图像中的样式的上下文的信息进行比较。然后模型想象或幻觉应该有什么颜色。

效果很好

在这里，该模型已经对目标图像/地面实况中的颜色进行了几乎完美的预测。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

该模型对目标图像/地面实况中的颜色进行了非常好的和令人信服的预测。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

这个预测给我留下了深刻印象，看起来很真实。模型无法预测某些颜色，伪装颜色是错误的，但它们看起来很有说服力，孤立地面真相/目标。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

另一个非常令人信服的颜色预测。在我看来，即使颜色不完全匹配，预测看起来与现实/目标一样逼真：

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

这里的头发颜色是斑点，肤色非常接近。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

该模型预测了猫毛中深色丰富的颜色，眼睛的颜色是斑点。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

几乎完美的颜色预测，除了金属屋顶的金色色调。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

有趣的结果

尽管颜色显然是错误的，但如果你看不到真相那么预测就会令人信服。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

有些颜色显然是错误的，但模型将如何知道花朵的颜色是什么？如果你看不到真相，那么预测看起来会很有说服力。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

结果不佳

论文和文章往往表现出最好的结果，而不是模特表现不佳的地方。重要的是突出模型的局限性和/或不能很好地概括的地方。
一般来说，模型对绿色的预测缺乏鲜明的色调，这在这里非常明显。果岭非常柔和，通常缺乏对比度。如果黄色/绿色更饱和，则预测会更接近。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

这里的模型并没有真正创造任何颜色，预测类似于棕褐色图像。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

结论

这种方法对于测试的大部分验证图像非常有效，几乎总能产生令人信服的彩色图像。
对于一个人来说，使这些颜色变得非常缓慢，在我看来，我会说这比人类的尝试更好，并且可能被视为超人类。

还有一些结果

以下是模型和基础事实的一些预测，以便与介绍性示例进行比较。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

预测几乎与目标无法区分：

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

同样，预测几乎与目标无法区分：

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

虽然预测的颜色与目标不同，但它非常有说服力：

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

该模型已经确定图像中有水，但不知道水可能有绿色而不是蓝色。天空几乎是完美的。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

狐狸非常接近目标/地面事实，天空和草地没有正确的颜色，尽管它们的颜色看起来很有说服力。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

狮子的毛皮颜色的变化并不完美，草/植物的绿色缺乏对比度，悬停生成的预测图像非常有说服力。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

一个复杂的图像，许多对象已经正确地着色，其他对象令人信服：

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

另一个令人信服的预测原始主题/目标中的颜色。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

除了一些小细节，例如浮标和小木屋的颜色，大部分颜色都被正确预测。

此模型将Greyscale 1通道图像增强为3通道彩色图像的示例。

技术细节

该模型是使用Fastai动态U-Net学习器https://docs.fast.ai/vision.models.unet.html创建的。
这使用Fastai软件库，PyTorch深度学习平台和CUDA并行计算API。
所有训练均在具有12GB RAM的Nvidia Tesla K80 GPU上进行，并且从开始到结束不到12小时，逐步调整大小。

Fastai

感谢Fastai团队，没有您的课程和软件库，我怀疑我是否能够进行这些实验并了解这些技术。