Skip to main content

基于深度学习的图像修复入门

Created on February 12|Last edited on July 12

导言




Select runs that logged labels to visualize data here.
Select runs that logged images to visualize data here.
Select runs that logged predictions to visualize data here.
Run set 2


什么是图像修复?

图像修复是重建图像的损坏/缺失部分的艺术,可以轻松地扩展到视频。由于图像修复的发展,大量用例已经出现。

(图像修复结果来自NVIDIA网络游乐场)

想象一下,你有一张小时候和祖父母一起拍的、你最喜欢的老照片,但是由于某些原因,这张照片的某些部分被损坏了。这是您最不想看到的,因为这张照片对你来说是如此特别。此时,图像修复可以说是你的救星。

对于可能没有预算来聘请艺术家来修复损坏画作的博物馆而言,图像修复可能会非常有用。

现在,想一下您最喜欢的照片编辑器。拥有图像修复功能会很酷,不是吗?

图像修复也可以扩展到视频(毕竟视频就是一系列的图像帧组成的)。由于过度压缩,视频的某些部分有时很可能会损坏。这种情况下,现代图像修复技术也能够优雅地处理此问题。

人工智能图像修复工具的主要目标是将缺失的部分填满,使其在视觉上和语义上都具有吸引力。但我们完全可以承认,这确实是一项具有挑战性的任务。

现在,我们对图像修复的含义有了一定的了解(我们将在以后进行更正式的定义),并了解了其中的一些用例,让我们切换一下角度,讨论一些用于修补图像的常用技术(剧透预警:经典计算机视觉)。

图像修复:传统方式

整个计算机视觉世界没有深度学习。在单帧检测器(SSD)出现之前,我们仍然可以进行对象检测(尽管其精度远不能达到SSD所能提供的能力)。类似地,有一些经典的计算机视觉技术可以进行图像修复。在本节中,我们将讨论其中两个。首先,让我们介绍一下这些技术所基于的中心主题——纹理合成(texture synthesis)或补丁合成(patch synthesis)。

为了修复图像中的特定缺失区域,他们从给定图像的周围区域借用了像素。值得注意的是,这些技术擅长在图像中修复背景,但不能推广到以下情况:

  • 周围区域可能没有合适的信息(读取像素)来填充缺失的部分。
  • 缺少的区域需要修复系统来推断可能存在的对象的属性。

如果我们从非常细的层次来看,图像修复只不过是恢复丢失的像素值。因此,我们可能会问自己——为什么我们不能把它看作是另一个缺失值填充的问题?嗯,图像并不是简单的像素值的任何随机集合,而是像素值的空间集合。因此,将图像修复任务仅仅视为缺失值填充问题是有点不合理的。我们稍后将回答以下问题——为什么不简单地使用CNN来预测丢失的像素?

现在,我们对图像修复的含义有了一定的了解(稍后我们将讨论更正式的定义),并了解了其中的一些用例,让我们切换一下角度,讨论一些用于修复图像的常用技术(剧透预警:经典计算机视觉)。

现在,介绍两种技术— -

  • Navier-Stokes方法 : 这种方法可以追溯到2001年(论文),并结合了流体力学和偏微分方程的概念。它基于这样一个事实,即图像中的边缘实际上应该是连续的。看下面的图形 -

image.png

  • 除了连续性约束(这不过是保留边缘特征的另一种说法)之外,作者还从需要修补的边缘周围区域提取颜色信息。
  • Fast marching方法:快速步进法:2004年,Alexandru Telea在论文中提出了这个想法。他提出以下建议:
  • 要估计丢失的像素,可以从像素附近获取像素的标准化加权总和。该邻域被一个边界参数化,一旦修补了一组像素,该边界就会更新。
  • 为了估计像素的颜色,使用了领域像素的梯度。

要了解这两种方法可以产生的结果,请参考这篇文章。现在,我们已经熟悉了进行图像修复的传统方式,让我们看看如何以现代方式(即深度学习)进行修复。

图像修复:现代方式

在这种方法中,我们训练了一个神经网络来预测图像的缺失部分,从而使预测在视觉和语义上都是一致的。让我们退后一步,思考一下我们(人类)如何进行图像修复。这将帮助我们制定基于深度学习的方法的基础。这也将有助于我们为图像修复任务形成问题陈述。

尝试重建图像中的缺失部分时,我们会利用对自身对世界的理解,并结合完成任务所需的上下文。这是一个我们很好地将某种上下文与全局理解相结合的例子。那么,我们可以将其运用到深度学习的模型中吗?我们可以试一试。

我们人类依赖于随着时间的推移获得的知识基础(对世界的了解)。当前的深度学习方法在任何意义上都远没有达到利用知识库的水平。但是我们肯定可以使用深度学习在图像中捕获空间上下文。卷积神经网络或CNN是用于处理具有已知网格状拓扑的数据的专用神经网络——例如,图像可以被视为像素的2D网格。这将是一种基于学习的方法,在该方法中,我们将训练基于深度CNN的架构来预测丢失的像素。

基于CIFA10数据集的简单图像修复模型

理解ML / DL概念的最好方法是实操。在本节中,我们将引导您完成深度图像修复,同时讨论其中的几个关键组成部分。我们首先需要一个数据集,最重要的是准备好它以适应目标任务。在讨论模型体系结构之前,放一个剧透——这项DL任务处在一个自我监督的学习环境中。

为什么选择一个简单的数据集?

由于修复是重建图像丢失或损坏部分的过程,因此我们可以获取任何图像数据集并为其添加人为损坏。对于此特定的DL任务,我们有大量的数据集可以使用。话虽如此,我们发现图像修补的实际应用是在高分辨率图像(例如512 x 512像素)上完成的。但是根据论文,要使一个像素受64个像素以外的内容的影响,它至少需要6层3×3卷积,其中膨胀因子为2。

因此,使用这种高分辨率图像不适合此处的目的。将ML / DL概念应用于玩具数据集是一种普遍做法。为了缩短计算资源并快速实施,我们将使用CIFAR10数据集。

数据准备

当然,任何DL任务的输入步骤都是数据准备。如前所述,在我们的案例中,我们需要对图像添加人为的损坏。这可以使用掩膜图像的标准图像处理思路来完成。由于它是在自我监督的学习环境中完成的,因此我们需要X和y(与X相同)数据对来训练我们的模型。这里X将是一批蒙版(masked)图像,而y将是原始/ground truth真实图像。

image.png

为了简化masking,我们首先假定丢失的部分是一个方孔。为了防止过度拟合此类伪像,我们将正方形的位置及其尺寸随机化。

使用这些方孔会大大限制模型在应用中的实用性。这是因为在现实生活中图像的损坏不只是方块形状。因此,受本文启发,我们将不规则的孔用作mask。我们使用OpenCV简单绘制了长度和粗细随机的线。

我们将使用Keras数据生成器以执行相同的操作。它将按照所需批次大小负责创建随机批次的X和y,将掩膜应用到X上并使其即时可用。对于高分辨率图像,使用数据生成器是唯一具有成本效益的选择。我们的数据生成器createAugment受到这个了不起的博客的启发。请好好读一下它。

class createAugment(keras.utils.Sequence):
 # Generates masked_image, masks, and target images for training
 def __init__(self, X, y, batch_size=32, dim=(32, 32),
	n_channels=3, shuffle=True):
	# Initialize the constructor
	self.batch_size = batch_size
	self.X = X
	self.y = y
	self.dim = dima
	self.n_channels = n_channels
	self.shuffle = shuffle
	self.on_epoch_end()

 def __len__(self):
	# Denotes the number of batches per epoch
	return int(np.floor(len(self.X) / self.batch_size))

 def __getitem__(self, index):
	# Generate one batch of data
	indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]
	# Generate data
	X_inputs, y_output = self.__data_generation(indexes)
	return X_inputs, y_output

 def on_epoch_end(self):
	# Updates indexes after each epoch
	self.indexes = np.arange(len(self.X))
	if self.shuffle:
		np.random.shuffle(self.indexes)

上面的代码块中的方法是不言自明的。我们来讨论一下专门针对我们的用例实施的方法——数据生成和createMask。顾名思义,此方法负责为给定批次大小的批次中的每个图像生成二进制掩码(binary masks)。它在白色背景上绘制长度和粗细随机的黑线。您可能会注意到,它与经过掩码处理的图像一起返回了mask。 为什么需要这个mask?我们很快就会看到。

Kera的model.fit需要输入和目标数据,在后台将其称为__getitem__。如果traingen是createAugment的一个实例,那么traingen [i]大致等效于traingen .__ getitem __(i),其中i的范围是0 to len(traingen)。此特殊方法在内部调用__data_generation,负责准备一批Masked_imagesMask_batchy_batch

def __data_generation(self, idxs):
	# Masked_images is a matrix of masked images used as input
	Masked_images = np.empty((self.batch_size, self.dim[0], self.dim[1], self.n_channels)) # Masked image

	# Mask_batch is a matrix of binary masks used as input
	Mask_batch = np.empty((self.batch_size, self.dim[0], self.dim[1], self.n_channels)) # Binary Masks
	
	# y_batch is a matrix of original images used for computing error from reconstructed image
	y_batch = np.empty((self.batch_size, self.dim[0], self.dim[1], self.n_channels)) # Original image
 
	## Iterate through random indexes
	for i, idx in enumerate(idxs):
		image_copy = self.X[idx].copy()
		## Get mask associated to that image
		masked_image, mask = self.__createMask(image_copy)

		## Append and scale down.
		Masked_images[i,] = masked_image/255
		Mask_batch[i,] = mask/255
		y_batch[i] = self.y[idx]/255

	return [Masked_images, Mask_batch], y_batch

架构

图像修复是大量图像生成问题中的一部分。修复的目的是填充丢失的像素。可以将其视为创建或修改像素,其中还包括诸如去模糊、去噪、伪像去除等任务。解决这些问题的方法通常依赖于自动编码器——一种经过训练的神经网络,可将其输入复制到其输出。它由一个学习描述输入的代码的编码器h = f(x)和一个产生重构的解码器r = g(h)r = g(f(x))组成。

普通卷积自动编码器

训练自动编码器的目的是为了重建输入,即g(f(x))= x,但这不是唯一的动因。我们希望自动编码器的训练会导致h具有区分性。人们已经注意到,如果没有对自动编码器进行仔细的训练,那么它往往会“记忆”数据而不会学习任何有用的显着特征。

我们选择使用规范化(regularized)的自动编码器,而不是限制(浅网络)编码器和解码器的能力。通常,使用损失函数的目的,是鼓励模型学习复制输入内容之外的其他属性。这些其他属性可以包括表示的稀疏性,对噪声或缺失值的稳健性。这是图像修补可以从基于自动编码器的模型结构中受益的地方。让我们来建立一个。

image.png

为了设置基准模型,我们将使用普通CNN构建自动编码器。最好先建立一个简单的模型来设定基准,然后进行逐步改进,这始终是一个好习惯。如果您想复习有关自动编码器的概念,PyImageSearch的这篇文章是一个很好的参考资料。如前所述,我们的目标不是掌握复制,因此我们设计损失函数,以便模型学习填充缺失点。如前所述,我们的目标不是掌握复制,因此我们设计损失函数,以便模型学习填充缺失点。我们以mean_square_error作为损失开始,并以dice coefficient作为评估指标。

def dice_coef(y_true, y_pred):
	y_true_f = keras.backend.flatten(y_true)
	y_pred_f = keras.backend.flatten(y_pred)
	intersection = keras.backend.sum(y_true_f * y_pred_f)
	return (2. * intersection) / (keras.backend.sum(y_true_f + y_pred_f))

对于图像分割、图像修复等任务,由于颜色类别高度不平衡,像素级别的精度并不是一个好的指标。尽管很容易解释,但准确性得分经常会误导人。两种常用的替代方法是交并比IoU和Dice系数。它们两者从某种意义上来说都是相似的——目标都是最大化预测像素与地表真实像素之间的重叠面积除以其并集。这里解释得很好,你可以看一下

看看模型是如何在不同epoch或step上学习填补缺失的洞洞的,这难道不是很有趣吗?

我们实施了一个简单的PredictionLogger回调演示——在每个epoch完成后,在大小为32的同一测试批次上调用model.predict()。使用wandb.log(),我们可以轻松地记录masked图像,mask,预测和ground truth图像。图1是此回调的结果。这是实现此功能的完整回调–

class PredictionLogger(tf.keras.callbacks.Callback):
	def __init__(self):
		super(PredictionLogger, self).__init__()
   
# The callback will be executed after an epoch is completed
	def on_epoch_end(self, logs, epoch):
		# Pick a batch, and sample the masked images, masks, and the labels
		sample_idx = 54
		[masked_images, masks], sample_labels = testgen[sample_idx]  

		# Initialize empty lists store intermediate results
		m_images = []
		binary_masks = []
		predictions = []
		labels = []
       
        # Iterate over the batch
		for i in range(32):
			 # Our inpainting model accepts masked imaged and masks as its inputs,
			 # then use perform inference
			 inputs = [B]
			 impainted_image = model.predict(inputs)

			 # Append the results to the respective lists
			 m_images.append(masked_images[i])
			 binary_masks.append(masks[i])
			 predictions.append(impainted_image.reshape(impainted_image.shape[1:]))
			 labels.append(sample_labels[i])

		# Log the results on wandb run page and voila!
		wandb.log({"masked_images": [wandb.Image(m_image)
		                     for m_image in m_images]})
		wandb.log({"masks": [wandb.Image(mask)
		                     for mask in binary_masks]})
		wandb.log({"predictions": [wandb.Image(inpainted_image)
		                     for inpainted_image in predictions]})
		wandb.log({"labels": [wandb.Image(label)
		                     for label in labels]})



Run set
0


部分卷积

现在,我们将讨论使用部分卷积进行不规则孔的图像修复(Image Inpainting for Irregular Holes Using Partial Convolutions),以替代普通CNN。部分卷积被提议用填充丢失的数据,例如图像中的孔。原始公式如下——假设X是当前滑动(卷积)窗口的特征值,而M是相应的二进制掩码。假设孔用0表示,非孔用1表示。在数学上,部分卷积可

image.png

比例因子sum(1)/ sum(M)会应用适当的比例来适应有效(未掩膜的)输入的量的变化。在每次部分卷积运算之后,我们按如下方式更新掩码:如果卷积能够以至少一个有效的输入(特征)值来调节其输出,则我们将该位置标记为有效。可以表示为

image.png

如果输入包含任何有效像素,则通过多层部分卷积,任何掩码最终都将全部为1。为了在图像修复任务中用部分卷积层替换普通CNN,我们需要同样的操作。

不幸的是,由于没有在TensorFlow和Pytorch中正式实现,因此我们必须自己实现该自定义层。关于如何构建自定义层的TensorFlow教程是一个很好的开始。幸运的是,我在这里可以找到部分卷积的Keras实现。该代码库使用TF 1.x作为Keras后端,我将其升级为使用TF2.x。我们已经为该博客文章提供了此升级版本以及GitHub贴文。在这里找到PConv2D。

让我们以代码实现该模型,并在CIFAR 10数据集上对其进行训练。我们实现了一个inpaintingModel类。要构建模型,您需要调用prepare_model()方法。

.

def prepare_model(self, input_size=(32,32,3)):
	input_image = keras.layers.Input(input_size)
	input_mask = keras.layers.Input(input_size)

	conv1, mask1, conv2, mask2 = self.__encoder_layer(32, input_image, input_mask)
	conv3, mask3, conv4, mask4 = self.__encoder_layer(64, conv2, mask2)
	conv5, mask5, conv6, mask6 = self.__encoder_layer(128, conv4, mask4)
	conv7, mask7, conv8, mask8 = self.__encoder_layer(256, conv6, mask6)

	conv9, mask9, conv10, mask10 = self.__decoder_layer(256, 128, conv8, mask8, conv7, mask7)
	conv11, mask11, conv12, mask12 = self.__decoder_layer(128, 64, conv10, mask10, conv5, mask5)
	conv13, mask13, conv14, mask14 = self.__decoder_layer(64, 32, conv12, mask12, conv3, mask3)
	conv15, mask15, conv16, mask16 = self.__decoder_layer(32, 3, conv14, mask14, conv1, mask1)

	outputs = keras.layers.Conv2D(3, (3, 3), activation='sigmoid', padding='same')(conv16)

	return keras.models.Model(inputs=[input_image, input_mask], outputs=[outputs])

由于它是自动编码器,因此该架构包含两个组件——编码器和解码器,这些我们已经讨论过了。为了重新使用编码器和解码器的卷积残差块,我们构建了两个简单的效用函数coder_layercoder_layer

def __encoder_layer(self, filters, in_layer, in_mask):
	conv1, mask1 = PConv2D(32, (3,3), strides=1, padding='same')([in_layer, in_mask])
	conv1 = keras.activations.relu(conv1)

	conv2, mask2 = PConv2D(32, (3,3), strides=2, padding='same')([conv1, mask1])
	conv2 = keras.layers.BatchNormalization()(conv2, training=True)
	conv2 = keras.activations.relu(conv2)

	return conv1, mask1, conv2, mask2

def __decoder_layer(self, filter1, filter2, in_img, in_mask, share_img, share_mask):
	up_img = keras.layers.UpSampling2D(size=(2,2))(in_img)
	up_mask = keras.layers.UpSampling2D(size=(2,2))(in_mask)
	concat_img = keras.layers.Concatenate(axis=3)([share_img, up_img])
	concat_mask = keras.layers.Concatenate(axis=3)([share_mask, up_mask])

	conv1, mask1 = PConv2D(filter1, (3,3), padding='same')([concat_img, concat_mask])
	conv1 = keras.activations.relu(conv1)

	conv2, mask2 = PConv2D(filter2, (3,3), padding='same')([conv1, mask1])
	conv2 = keras.layers.BatchNormalization()(conv2)
	conv2 = keras.activations.relu(conv2)

	return conv1, mask1, conv2, mask2

Autoencoder实现的本质在于Upsampling2DConcatenate层。替代方法是使用Conv2DTranspose层。

我们使用具有默认参数的Adam优化器编译了模型,以mean_square_error作为损失函数, dice_coef作为指标。我们使用model.fit()训练了模型,并使用WandbCallback和PredictionLogger回调记录了结果。




Run set


结论

我们以该主题外的其他一些讨论作为结束语,包括它与自我监督学习的关系以及进行图像修复的最新方法。

图像修复模型的一个非常有趣的特性是它能够在某种程度上理解图像。与NLP非常相似,在NLP中,我们使用嵌入来理解单词之间的语义关系,并将这些嵌入用于诸如文本分类之类的下游任务。

这里的前提是,当您开始用有吸引力的语义和视觉填充图像的缺失部分时,您就开始理解该图像。这更像是自我监督学习,在没有任何显式标签的情况下,您可以利用输入数据中存在的隐式标签。

这特别有趣,因为我们可以在计算机视觉任务中使用图像修复模型的知识,就像我们将嵌入用于NLP任务一样。要了解更多有关此内容的信息,我强烈推荐Jeremy Howard撰写的这篇出色的文章

到目前为止,我们仅使用逐像素比较(pixel-wise comparison)作为损失函数。这通常导致我们的模型学习非常僵化和不太丰富的特征表示。 Charles等人提出了一个非常有趣但简单的想法,即近似精确匹配(approximate exact matching)。在这份报告中。根据他们的研究,如果我们将图像的像素值移动一个小的常数,不会使图像在视觉上与原始形式有很大不同。因此,他们在逐像素比较中增加了一个附加项,以体现这一思想。

对我们网络的另一个有趣的想法,是使它能够参与图像中遥远空间位置的相关特征补丁。在Generative Image Inpainting with Contextual Attention(使用内容感知进行图像修复)这篇文章中, Jiahui 等人引入了内容感知attention的思想,该思想允许网络在训练期间显式地利用邻近的图像特征作为参考。

感谢你读完这篇文章。图像修复是一项非常有趣的计算机视觉任务,我们希望本文能让你对这个主题有一个很好的了解。如果你有任何反馈,请随时通过Twitter(AyushSayak)联系我们。谢谢:)


Iterate on AI agents and models faster. Try Weights & Biases today.