开发者社区> 【方向】> 正文
阿里云
为了无法计算的价值
打开APP
阿里云APP内打开

老旧黑白片修复机——使用卷积神经网络图像自动着色实战(附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地点数据集中的一个子集,该子数据集包含地点、景观和建筑物。

# Download and unzip (2.2GB)
!wget http://data.csail.mit.edu/places/places205/testSetPlaces205_resize.tar.gz
!tar -xzf testSetPlaces205_resize.tar.gz
# Move data into training and validation directories
import os
os.makedirs('images/train/class/', exist_ok=True) # 40,000 images
os.makedirs('images/val/class/', exist_ok=True)   #  1,000 images
for i, file in enumerate(os.listdir('testSet_resize')):
  if i < 1000: # first 1000 will be val
    os.rename('testSet_resize/' + file, 'images/val/class/' + file)
  else: # others will be val
    os.rename('testSet_resize/' + file, 'images/train/class/' + file)
# Make sure the images are there
from IPython.display import Image, display
display(Image(filename='images/val/class/84b3ccd8209a4db1835988d28adfed4c.jpg'))

1

工具

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

# Download and import libraries
!pip install torch torchvision matplotlib numpy scikit-image pillow==4.1.1
# For plotting
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
# For conversion
from skimage.color import lab2rgb, rgb2lab, rgb2gray
from skimage import io
# For everything
import torch
import torch.nn as nn
import torch.nn.functional as F
# For our model
import torchvision.models as models
from torchvision import datasets, transforms
# For utilities
import os, shutil, time
# Check if GPU is available
use_gpu = torch.cuda.is_available()

模型

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

2


       现在,在代码中定义后续的网络模型,将从网络的后半部分开始,即上采样层:
class ColorizationNet(nn.Module):
  def __init__(self, input_size=128):
    super(ColorizationNet, self).__init__()
    MIDLEVEL_FEATURE_SIZE = 128

    ## First half: ResNet
    resnet = models.resnet18(num_classes=365) 
    # Change first conv layer to accept single-channel (grayscale) input
    resnet.conv1.weight = nn.Parameter(resnet.conv1.weight.sum(dim=1).unsqueeze(1)) 
    # Extract midlevel features from ResNet-gray
    self.midlevel_resnet = nn.Sequential(*list(resnet.children())[0:6])

    ## Second half: Upsampling
    self.upsample = nn.Sequential(     
      nn.Conv2d(MIDLEVEL_FEATURE_SIZE, 128, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(128),
      nn.ReLU(),
      nn.Upsample(scale_factor=2),
      nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(64),
      nn.ReLU(),
      nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(64),
      nn.ReLU(),
      nn.Upsample(scale_factor=2),
      nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(32),
      nn.ReLU(),
      nn.Conv2d(32, 2, kernel_size=3, stride=1, padding=1),
      nn.Upsample(scale_factor=2)
    )

  def forward(self, input):

    # Pass input through ResNet-gray to extract features
    midlevel_features = self.midlevel_resnet(input)

    # Upsample to get colors
    output = self.upsample(midlevel_features)
    return output

       现在通过下面的代码创建整个模型:

model = ColorizationNet()

训练

损失函数

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

criterion = nn.MSELoss()

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

优化

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

optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=0.0)

加载数据

       使用torchtext来加载数据,由于我们需要LAB空间中的图像,所以首先必须定义一个自定义数据加载器(dataloader)来转换图像。

class GrayscaleImageFolder(datasets.ImageFolder):
  '''Custom images folder, which converts images to grayscale before loading'''
  def __getitem__(self, index):
    path, target = self.imgs[index]
    img = self.loader(path)
    if self.transform is not None:
      img_original = self.transform(img)
      img_original = np.asarray(img_original)
      img_lab = rgb2lab(img_original)
      img_lab = (img_lab + 128) / 255
      img_ab = img_lab[:, :, 1:3]
      img_ab = torch.from_numpy(img_ab.transpose((2, 0, 1))).float()
      img_original = rgb2gray(img_original)
      img_original = torch.from_numpy(img_original).unsqueeze(0).float()
    if self.target_transform is not None:
      target = self.target_transform(target)
    return img_original, img_ab, target

       接下来,对训练数据和验证数据定义变换。

# Training
train_transforms = transforms.Compose([transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip()])
train_imagefolder = GrayscaleImageFolder('images/train', train_transforms)
train_loader = torch.utils.data.DataLoader(train_imagefolder, batch_size=64, shuffle=True)

# Validation 
val_transforms = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(224)])
val_imagefolder = GrayscaleImageFolder('images/val' , val_transforms)
val_loader = torch.utils.data.DataLoader(val_imagefolder, batch_size=64, shuffle=False)

辅助函数

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

class AverageMeter(object):
  '''A handy class from the PyTorch ImageNet tutorial''' 
  def __init__(self):
    self.reset()
  def reset(self):
    self.val, self.avg, self.sum, self.count = 0, 0, 0, 0
  def update(self, val, n=1):
    self.val = val
    self.sum += val * n
    self.count += n
    self.avg = self.sum / self.count

def to_rgb(grayscale_input, ab_input, save_path=None, save_name=None):
  '''Show/save rgb image from grayscale and ab channels
     Input save_path in the form {'grayscale': '/path/', 'colorized': '/path/'}'''
  plt.clf() # clear matplotlib 
  color_image = torch.cat((grayscale_input, ab_input), 0).numpy() # combine channels
  color_image = color_image.transpose((1, 2, 0))  # rescale for matplotlib
  color_image[:, :, 0:1] = color_image[:, :, 0:1] * 100
  color_image[:, :, 1:3] = color_image[:, :, 1:3] * 255 - 128   
  color_image = lab2rgb(color_image.astype(np.float64))
  grayscale_input = grayscale_input.squeeze().numpy()
  if save_path is not None and save_name is not None: 
    plt.imsave(arr=grayscale_input, fname='{}{}'.format(save_path['grayscale'], save_name), cmap='gray')
    plt.imsave(arr=color_image, fname='{}{}'.format(save_path['colorized'], save_name))

验证

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

def validate(val_loader, model, criterion, save_images, epoch):
  model.eval()

  # Prepare value counters and timers
  batch_time, data_time, losses = AverageMeter(), AverageMeter(), AverageMeter()

  end = time.time()
  already_saved_images = False
  for i, (input_gray, input_ab, target) in enumerate(val_loader):
    data_time.update(time.time() - end)

    # Use GPU
    if use_gpu: input_gray, input_ab, target = input_gray.cuda(), input_ab.cuda(), target.cuda()

    # Run model and record loss
    output_ab = model(input_gray) # throw away class predictions
    loss = criterion(output_ab, input_ab)
    losses.update(loss.item(), input_gray.size(0))

    # Save images to file
    if save_images and not already_saved_images:
      already_saved_images = True
      for j in range(min(len(output_ab), 10)): # save at most 5 images
        save_path = {'grayscale': 'outputs/gray/', 'colorized': 'outputs/color/'}
        save_name = 'img-{}-epoch-{}.jpg'.format(i * val_loader.batch_size + j, epoch)
        to_rgb(input_gray[j].cpu(), ab_input=output_ab[j].detach().cpu(), save_path=save_path, save_name=save_name)

    # Record time to do forward passes and save images
    batch_time.update(time.time() - end)
    end = time.time()

    # Print model accuracy -- in the code below, val refers to both value and validation
    if i % 25 == 0:
      print('Validate: [{0}/{1}]\t'
            'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
            'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(
             i, len(val_loader), batch_time=batch_time, loss=losses))

  print('Finished validation.')
  return losses.avg

训练

       在训练过程中,使用loss.backward()运行模型并进行反向传播过程。我们首先定义了一个训练一个epoch的函数:

def train(train_loader, model, criterion, optimizer, epoch):
  print('Starting training epoch {}'.format(epoch))
  model.train()
  
  # Prepare value counters and timers
  batch_time, data_time, losses = AverageMeter(), AverageMeter(), AverageMeter()

  end = time.time()
  for i, (input_gray, input_ab, target) in enumerate(train_loader):
    
    # Use GPU if available
    if use_gpu: input_gray, input_ab, target = input_gray.cuda(), input_ab.cuda(), target.cuda()

    # Record time to load data (above)
    data_time.update(time.time() - end)

    # Run forward pass
    output_ab = model(input_gray) 
    loss = criterion(output_ab, input_ab) 
    losses.update(loss.item(), input_gray.size(0))

    # Compute gradient and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Record time to do forward and backward passes
    batch_time.update(time.time() - end)
    end = time.time()

    # Print model accuracy -- in the code below, val refers to value, not validation
    if i % 25 == 0:
      print('Epoch: [{0}][{1}/{2}]\t'
            'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
            'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
            'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(
              epoch, i, len(train_loader), batch_time=batch_time,
             data_time=data_time, loss=losses)) 

  print('Finished training epoch {}'.format(epoch))

       接下来,我们定义一个循环训练函数,即训练100个epoch:

# Move model and loss function to GPU
if use_gpu: 
  criterion = criterion.cuda()
  model = model.cuda()
# Make folders and set parameters
os.makedirs('outputs/color', exist_ok=True)
os.makedirs('outputs/gray', exist_ok=True)
os.makedirs('checkpoints', exist_ok=True)
save_images = True
best_losses = 1e10
epochs = 100
# Train model
for epoch in range(epochs):
  # Train for one epoch, then validate
  train(train_loader, model, criterion, optimizer, epoch)
  with torch.no_grad():
    losses = validate(val_loader, model, criterion, save_images, epoch)
  # Save checkpoint and replace old best model if current model is better
  if losses < best_losses:
    best_losses = losses
    torch.save(model.state_dict(), 'checkpoints/model-epoch-{}-losses-{:.3f}.pth'.format(epoch+1,losses))
Starting training epoch 0 ...

预训练模型

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

# Download pretrained model
!wget https://www.dropbox.com/s/kz76e7gv2ivmu8p/model-epoch-93.pth
#https://www.dropbox.com/s/9j9rvaw2fo1osyj/model-epoch-67.pth
# Load model
pretrained = torch.load('model-epoch-93.pth', map_location=lambda storage, loc: storage)
model.load_state_dict(pretrained)
# Validate
save_images = True
with torch.no_grad():
  validate(val_loader, model, criterion, save_images, 0)
Validate: [0/16]    Time 10.628 (10.628)    Loss 0.0030 (0.0030)
Validate: [16/16]   Time  0.328 ( 0.523)    Loss 0.0029 (0.0029)  

结果

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

# Show images 
import matplotlib.image as mpimg
image_pairs = [('outputs/color/img-2-epoch-0.jpg', 'outputs/gray/img-2-epoch-0.jpg'),
               ('outputs/color/img-7-epoch-0.jpg', 'outputs/gray/img-7-epoch-0.jpg')]
for c, g in image_pairs:
  color = mpimg.imread(c)
  gray  = mpimg.imread(g)
  f, axarr = plt.subplots(1, 2)
  f.set_size_inches(15, 15)
  axarr[0].imshow(gray, cmap='gray')
  axarr[1].imshow(color)
  axarr[0].axis('off'), axarr[1].axis('off')
  plt.show()

00


01

结论

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

数十款阿里云产品限时折扣中,赶紧点击领劵开始云上实践吧!

作者信息

Luke Melas,专研于数学与计算机科学
个人主页:https://lukemelas.github.io/image-colorization.html
本文由阿里云云栖社区组织翻译。
文章原标题《Image Colorization with Convolutional Neural Networks》,译者:海棠,审校:[Uncle_LLD]()。
文章为简译,更为详细的内容,请查看原文

版权声明:本文内容由阿里云实名注册用户自发贡献,版权归原作者所有,阿里云开发者社区不拥有其著作权,亦不承担相应法律责任。具体规则请查看《阿里云开发者社区用户服务协议》和《阿里云开发者社区知识产权保护指引》。如果您发现本社区中有涉嫌抄袭的内容,填写侵权投诉表单进行举报,一经查实,本社区将立刻删除涉嫌侵权内容。

相关文章
CentOS7 yum使用阿里源安装tomcat无欢迎页问题解决办法
解决过程 在阿里云服务器CentOS7上通过命令安装上了tomcat: yum -y install tomcat 接着执行命令,启动tomcat服务: systemctl start tomcat.service 但是当我在浏览器中输入ip:port后却没有反应,然后报404错误。
1871 0
【直播预告】:Java Spring Boot开发实战系列课程【第14讲】:Spring Boot 2.0实战MyBatis连接池阿里Druid与SQL性能监控
阿里开源数据库连接池组件Druid非常强大,,本次课程一起学习如何在最新的Java Spring Boot 2.0和MyBatis系统 中集成阿里开源的连接池Druid,以及SQL性能监控,生产环境必备利器。
66349 0
14.AndFix热修复的使用和源码分析(客户端修复逻辑)
使用方法: 1.下载AndFix源码:https://github.com/alibaba/AndFix 2.生成两个apk文件,一个是含有bug的old.
1057 0
《连载 | 物联网框架ServerSuperIO教程》- 17.集成Golden实时数据库,高并发保存测点数据。附:3.4 发布与版本更新说明。
1.C#跨平台物联网通讯框架ServerSuperIO(SSIO)介绍 《连载 | 物联网框架ServerSuperIO教程》1.4种通讯模式机制。 《连载 | 物联网框架ServerSuperIO教程》2.服务实例的配置参数说明 《连载 | 物联网框架ServerSuperIO教程》- 3.设备驱动介绍 《连载 | 物联网框架ServerSuperIO教程》-4.如开发一套设备驱动,同时支持串口和网络通讯。
857 0
[Spring实战系列](14)Bean的自动检测
版权声明:本文为博主原创文章,未经博主允许不得转载。 https://blog.csdn.net/SunnyYoona/article/details/50650102 即使有助于完全消除Spring注解中的和元素,但是还是不能完全消除,仍然需要使用元素显示定义Bean。
982 0
+关注
【方向】
欢迎各位对内容方向及质量提需求,我们尽量满足,将国外优质的内容呈现给大家!
文章
问答
文章排行榜
最热
最新
相关电子书
更多
140-弱监督机器学...1506573734.pdf
立即下载
阿里巴巴代码缺陷检测探索与实践
立即下载
利用CNN实现无需联网的图像识别
立即下载