卷积神经网络去马赛克（老旧黑白片修复机）

卷积神经网络去马赛克（老旧黑白片修复机）现在，在代码中定义后续的网络模型，将从网络的后半部分开始，即上采样层：具体来说，模型采用的是迁移学习的方法，基础是ResNet-18模型，ResNet-18网络具有18层结构以及剩余连接的图像分类网络层。我们修改了该网络的第一层，以便它接受灰度输入而不是彩色输入，并且切断了第六层后面的网络结构：本文使用PyTorch构建和训练搭建的模型。此外，我们还了使用torchvision工具，该工具在PyTorch中处理图像和视频时很有用，以及使用了scikit-learn工具，用于在RGB和LAB颜色空间之间进行转换。模型模型采用卷积神经网络构建而成，与传统的卷积神经网络模型类似，首先应用一些卷积层从图像中提取特征，然后将反卷积层应用于高级（增加空间分辨率）特征。

摘要： 照片承载了很多人在某个时刻的记忆，尤其是一些老旧的黑白照片，尘封于脑海之中，随着时间的流逝，记忆中对当时颜色的印象也会慢慢消散，这确实有些可惜。技术的发展会解决一些现有的难题，深度学习恰好能够解决这个问题。

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

PS：不仅能够对旧图像进行着色，还可以对视频（每次对视频进行一帧处理）进行着色哦！闲话少叙，下面直接进入正题吧。

工具

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

模型

模型采用卷积神经网络构建而成，与传统的卷积神经网络模型类似，首先应用一些卷积层从图像中提取特征，然后将反卷积层应用于高级（增加空间分辨率）特征。

具体来说，模型采用的是迁移学习的方法，基础是ResNet-18模型，ResNet-18网络具有18层结构以及剩余连接的图像分类网络层。我们修改了该网络的第一层，以便它接受灰度输入而不是彩色输入，并且切断了第六层后面的网络结构：

现在，在代码中定义后续的网络模型，将从网络的后半部分开始，即上采样层：

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

model = ColorizationNet()

训练

损失函数

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

criterion = nn.MSELoss()

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

优化

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

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

加载数据

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

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

# 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图像。

验证

在验证过程中，使用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()

结论

在这篇文章中，使用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等，说不定下一个风投就是你哦！

作者信息

Luke Melas，专研于数学与计算机科学

本文由阿里云云栖社区组织翻译。

文章原标题《Image Colorization with Convolutional Neural Networks》，译者：海棠，审校：[Uncle_LLD]()。