如何捕获一只彩色卓别林?黑白照片AI上色教程很友好

访问flyai.club,一键创建你的人工智能项目

老照片的手动着色魔法

妈妈小时候已经有彩色照片了,不过那些照片,还是照相馆的人类手动上色的。

几十年之后,人们已经开始培育深度神经网络,来给老照片和老电影上色了。

来自哈佛大学的Luke Melas-Kyriazi (我叫他卢克吧) ,用自己训练的神经网络,把卓别林变成了彩色的卓别林,清新自然。

作为一只哈佛学霸,卢克还为钻研机器学习的小伙伴们写了一个基于PyTorch的教程。

虽然教程里的模型比给卓别林用的模型要简约一些,但效果也是不错了。

问题是什么?

卢克说,给黑白照片上色这个问题的难点在于,它是多模态的——与一幅灰度图像对应的合理彩色图像,并不唯一。

这并不是正确示范

传统模型需要输入许多额外信息,来辅助上色。

而深度神经网络,除了灰度图像之外,不需要任何额外输入,就可以完成上色。

在彩色图像里,每个像素包含三个值,即亮度饱和度以及色调

而灰度图像,并无饱和度色调可言,只有亮度一个值。

所以,模型要用一组数据,生成另外两足数据。换句话说,以灰度图像为起点,推断出对应的彩色图像。

为了简单,这里只做了256 x 256像素的图像上色。输出的数据量则是256 x 256 x 2。

关于颜色表示,卢克用的是LAB色彩空间,它跟RGB系统包含的信息是一样的。

但对程序猿来说,前者比较方便把亮度和其他两项分离开来。

数据也不难获得,卢克用了MIT Places数据集,中的一部分。内容就是校园里的一些地标和风景。然后转换成黑白图像,就可以了。以下为数据搬运代码——

1# Download and unzip (2.2GB)

2!wget http://data.csail.mit.edu/places/places205/testSetPlaces205_resize.tar.gz

3!tar -xzf testSetPlaces205_resize.tar.gz

1# Move data into training and validation directories

2importos

3os.makedirs('images/train/class/', exist_ok=True)# 40,000 images

4os.makedirs('images/val/class/', exist_ok=True)# 1,000 images

5fori, fileinenumerate(os.listdir('testSet_resize')):

6ifi

7os.rename('testSet_resize/'+ file,'images/val/class/'+ file)

8else:# others will be val

9os.rename('testSet_resize/'+ file,'images/train/class/'+ file)

1# Make sure the images are there

2fromIPython.displayimportImage, display

3display(Image(filename='images/val/class/84b3ccd8209a4db1835988d28adfed4c.jpg'))

好用的工具有哪些?

搭建模型和训练模型是在PyTorch里完成的。

还用了torchvishion,这是一套在PyTorch上处理图像和视频的工具。

另外,scikit-learn能完成图片在RGB和LAB色彩空间之间的转换。

1# Download and import libraries

2!pip install torch torchvision matplotlib numpy scikit-image pillow==4.1.1

1# For plotting

2importnumpyasnp

3importmatplotlib.pyplotasplt

4%matplotlib inline

5# For conversion

6fromskimage.colorimportlab2rgb, rgb2lab, rgb2gray

7fromskimageimportio

8# For everything

9importtorch

10importtorch.nnasnn

11importtorch.nn.functionalasF

12# For our model

13importtorchvision.modelsasmodels

14fromtorchvisionimportdatasets, transforms

15# For utilities

16importos, shutil, time

模型长什么样?

神经网络里面,第一部分是几层用来提取图像特征;第二部分是一些反卷积层 (Deconvolutional Layers) ,用来给那些特征增加分辨率。

具体来说,第一部分用的是ResNet-18,这是一个图像分类网络,有18层,以及一些残差连接 (Residual Connections) 。

给第一层做些修改,它就可以接受灰度图像输入了。然后把第6层之后的都去掉。

然后,用代码来定义一下这个模型。

从神经网络的第二部分(就是那些上采样层) 开始。

下一步,创建模型吧。

它是怎么训练的?

预测每个像素的色值,用的是回归 (Regression) 的方法。

损失函数 (Loss Function)

所以,用了一个均方误差 (MSE) 损失函数——让预测的色值与参考标准 (Ground Truth) 之间的距离平方最小化。

优化损失函数

这里是用Adam Optimizer优化的。

加载数据

用torchtext加载数据。首先定义一个专属的数据加载器 (DataLoader) ,来完成RGB到LAB空间的转换。

再来,就是定义训练数据验证数据的转换。

辅助函数 (Helper Function)

训练开始之前,要把辅助函数写好,来追踪训练损失,并把图像转回RGB形式。

验证

不用反向传播 (Back Propagation),直接用torch.no_grad() 跑模型。

训练

用loss.backward(),用上反向传播。写一下训练数据跑一遍 (one epoch) 用的函数。

然后,定义一个训练回路 (Training Loop) ,跑一百遍训练数据。从Epoch 0开始训练。

1# Move model and loss function to GPU

2ifuse_gpu:

3criterion = criterion.cuda()

4model = model.cuda()

1# Make folders and set parameters

2os.makedirs('outputs/color', exist_ok=True)

3os.makedirs('outputs/gray', exist_ok=True)

4os.makedirs('checkpoints', exist_ok=True)

5save_images =True

6best_losses =1e10

7epochs =100

训练结果什么样?

是时候看看修炼成果了,所以,复制一下这段代码。

效果还是很自然的,虽然生成的彩色图像不是那么明丽。

卢克说,问题是多模态的,所以损失函数还是值得推敲。

比如,一条灰色裙子可以是蓝色也可以是红色。如果模型选择的颜色和参考标准不同,就会受到严厉的惩罚。

这样一来,模型就会选择哪些不会被判为大错特错的颜色,而不太选择非常显眼明亮的颜色。

没时间怎么办?

卢克还把一只训练好的AI放了出来,不想从零开始训练的小伙伴们,也可以直接感受他的训练成果,只要用以下代码下载就好了。

1# Download pretrained model

2!wget https://www.dropbox.com/s/kz76e7gv2ivmu8p/model-epoch-93.pth

3#https://www.dropbox.com/s/9j9rvaw2fo1osyj/model-epoch-67.pth

1# Load model

2pretrained = torch.load('model-epoch-93.pth', map_location=lambdastorage, loc: storage)

3model.load_state_dict(pretrained)

彩色老电影?

如果想要更加有声有色的结局,就不能继续偷懒了。卢克希望大家沿着他精心铺就的路,走到更远的地方。

要替换当前的损失函数,可以参考Zhang et al. (2017):

https://richzhang.github.io/ideepcolor/

无监督学习的上色大法,可以参考Larsson et al. (2017):

http://people.cs.uchicago.edu/~larsson/color-proxy/

另外,可以做个手机应用,就像谷歌在I/O大会上发布的着色软件那样。

黑白电影,也可以自己去尝试,一帧一帧地上色。

这里有卓别林用到的完整代码

https://github.com/lukemelas/Automatic-Image-Colorization/

— End —

  • 发表于:
  • 原文链接https://kuaibao.qq.com/s/20180601A09ANU00?refer=cp_1026
  • 腾讯「腾讯云开发者社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。
  • 如有侵权,请联系 cloudcommunity@tencent.com 删除。

扫码关注腾讯云开发者

领取腾讯云代金券