计算机视觉:风格迁移

风格迁移

本节将介绍如何使用卷积神经网络,自动将一个图像中的风格应用在另一图像之上,即风格迁移(style transfer) (Gatys et al., 2016)。 这里我们需要两张输入图像:一张是内容图像,另一张是风格图像。 我们将使用神经网络修改内容图像,使其在风格上接近风格图像。 例如, 图中的内容图像为本书作者在西雅图郊区的雷尼尔山国家公园拍摄的风景照,而风格图像则是一幅主题为秋天橡树的油画。 最终输出的合成图像应用了风格图像的油画笔触让整体颜色更加鲜艳,同时保留了内容图像中物体主体的形状。

在这里插入图片描述



方法

首先,我们初始化合成图像,例如将其初始化为内容图像。 该合成图像是风格迁移过程中唯一需要更新的变量,即风格迁移所需迭代的模型参数。 然后,我们选择一个预训练的卷积神经网络来抽取图像的特征,其中的模型参数在训练中无须更新。 这个深度卷积神经网络凭借多个层逐级抽取图像的特征,我们可以选择其中某些层的输出作为内容特征或风格特征。 以下图为例,这里选取的预训练的神经网络含有3个卷积层,其中第二层输出内容特征,第一层和第三层输出风格特征。
在这里插入图片描述
接下来,我们通过前向传播(实线箭头方向)计算风格迁移的损失函数,并通过反向传播(虚线箭头方向)迭代模型参数,即不断更新合成图像。 风格迁移常用的损失函数由3部分组成:

  1. 内容损失使合成图像与内容图像在内容特征上接近;

  2. 风格损失使合成图像与风格图像在风格特征上接近;

  3. 全变分损失则有助于减少合成图像中的噪点。

最后,当模型训练结束时,我们输出风格迁移的模型参数,即得到最终的合成图像。

在下面,我们将通过代码来进一步了解风格迁移的技术细节。

阅读内容和风格图像

%matplotlib inline
import torch
import torchvision
from torch import nn
from d2l import torch as d2l

d2l.set_figsize()
content_img = d2l.Image.open('rainier.jpg')
d2l.plt.imshow(content_img);

在这里插入图片描述

style_img = d2l.Image.open('autumn-oak.jpg')
d2l.plt.imshow(style_img);

在这里插入图片描述

预处理和后处理

rgb_mean = torch.tensor([0.485, 0.456, 0.406])
rgb_std = torch.tensor([0.229, 0.224, 0.225])

def preprocess(img, image_shape):
    transforms = torchvision.transforms.Compose([
        torchvision.transforms.Resize(image_shape),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=rgb_mean, std=rgb_std)])
    return transforms(img).unsqueeze(0)

def postprocess(img):
    img = img[0].to(rgb_std.device)
    img = torch.clamp(img.permute(1, 2, 0) * rgb_std + rgb_mean, 0, 1)
    return torchvision.transforms.ToPILImage()(img.permute(2, 0, 1))

preprocess 函数接受一个输入图像 img 和一个 image_shape 参数作为输入。它使用 torchvision.transforms.Resize 将图像大小调整为指定的大小,使用 torchvision.transforms.ToTensor 将图像转换为张量,并使用 torchvision.transforms.Normalize 对张量进行归一化,其中的均值和标准差值分别由 rgb_mean 和 rgb_std 定义。最后,它使用 unsqueeze(0) 添加了一个额外的维度,以创建一个批次维度,然后返回预处理后的张量。

postprocess 函数接受一个张量 img 作为输入。它通过使用索引 [0] 移除批次维度,使用 .to(rgb_std.device) 将张量移动到 rgb_std 所在的设备上,然后使用 torch.clamp 和 * 操作对张量进行反归一化。torch.clamp 函数用于确保像素值在 0 和 1 之间,因为归一化可能导致某些值超出此范围。最后,它使用 torchvision.transforms.ToPILImage() 将张量转换回 PIL 图像,并使用 .permute(2, 0, 1) 对维度进行排列,以便通道成为最后一个维度,然后返回图像。

归一化是将数据转换为特定范围内的值的过程。在计算机视觉中,通常使用均值和标准差来对图像进行归一化。这是因为图像中的每个像素都可以视为一组数字,每个数字的值在 0 到 255 之间。因此,对图像进行归一化可以将像素值缩放到 0 到 1 之间,使它们更易于处理和比较。在本代码中,torchvision.transforms.Normalize 函数用于对图像进行归一化,其中 rgb_mean 和 rgb_std 定义了用于标准化图像的均值和标准差值。

反归一化是将数据从归一化范围内的值转换为原始值的过程。在计算机视觉中,反归一化通常用于将模型输出转换回原始图像的像素值范围。在本代码中,postprocess 函数中的 torch.clamp 和 * 操作用于对张量进行反归一化,以将像素值从归一化范围内的值转换为原始图像像素值。

抽取图像特征

pretrained_net = torchvision.models.vgg19(pretrained=True)
print(pretrained_net)

结构为:

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (17): ReLU(inplace=True)
    (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (24): ReLU(inplace=True)
    (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (26): ReLU(inplace=True)
    (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (31): ReLU(inplace=True)
    (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (33): ReLU(inplace=True)
    (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (35): ReLU(inplace=True)
    (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
 (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

为了抽取图像的内容特征和风格特征,我们可以选择VGG网络中某些层的输出。 一般来说,越靠近输入层,越容易抽取图像的细节信息反之,则越容易抽取图像的全局信息。 为了避免合成图像过多保留内容图像的细节,我们选择VGG较靠近输出的层,即内容层,来输出图像的内容特征。 我们还从VGG中选择不同层的输出来匹配局部和全局的风格,这些图层也称为风格层。 VGG网络使用了5个卷积块。 实验中,我们选择第四卷积块的最后一个卷积层作为内容层选择每个卷积块的第一个卷积层作为风格层。 这些层的索引可以通过打印pretrained_net实例获取。

style_layers, content_layers = [0, 5, 10, 19, 28], [25]

使用VGG层抽取特征时,我们只需要用到从输入层到最靠近输出层的内容层或风格层之间的所有层。 下面构建一个新的网络net,它只保留需要用到的VGG的所有层。

net = nn.Sequential(*[pretrained_net.features[i] for i in
                      range(max(content_layers + style_layers) + 1)])

这些层的选择是由 content_layers 和 style_layers 两个列表决定的。content_layers 列表包含用于内容提取的层的索引,而 style_layers 列表包含用于风格提取的层的索引。这些层的索引是根据 VGG19 的结构确定的,因为在该模型的特定层中提取的特征已被证明对于内容和风格的捕捉非常有用。

给定输入X,如果我们简单地调用前向传播net(X),只能获得最后一层的输出。 由于我们还需要中间层的输出,因此这里我们逐层计算,并保留内容层和风格层的输出。

def extract_features(X, content_layers, style_layers):
    contents = []
    styles = []
    for i in range(len(net)):
        X = net[i](X)
        if i in style_layers:
            styles.append(X)
        if i in content_layers:
            contents.append(X)
    return contents, styles

这段代码定义了一个函数 extract_features,它接受输入图像 X,以及用于内容提取和风格提取的层索引列表 content_layers 和 style_layers。

函数中的 for 循环逐层遍历了神经网络 net,并应用了输入图像 X,以获取每一层的输出。如果当前层的索引在 style_layers 中,则将该层的输出添加到 styles 列表中。如果当前层的索引在 content_layers 中,则将该层的输出添加到 contents 列表中。最后,函数返回 contents 和 styles 两个列表,分别包含了用于内容和风格提取的层的输出。

这个函数的目的是使用神经网络 net 中的层来提取输入图像的内容和风格特征。特别地,contents 列表包含了用于表示输入图像内容的特征,而 styles 列表包含了用于表示输入图像风格的特征。这些特征将用于计算损失函数,进而用于训练风格迁移模型。

下面定义两个函数:get_contents函数对内容图像抽取内容特征; get_styles函数对风格图像抽取风格特征。 因为在训练时无须改变预训练的VGG的模型参数,所以我们可以在训练开始之前就提取出内容特征和风格特征。 由于合成图像是风格迁移所需迭代的模型参数,我们只能在训练过程中通过调用extract_features函数来抽取合成图像的内容特征和风格特征。

def get_contents(image_shape, device):
    content_X = preprocess(content_img, image_shape).to(device)
    contents_Y, _ = extract_features(content_X, content_layers, style_layers)
    return content_X, contents_Y

def get_styles(image_shape, device):
    style_X = preprocess(style_img, image_shape).to(device)
    _, styles_Y = extract_features(style_X, content_layers, style_layers)
    return style_X, styles_Y

损失函数

下面我们来描述风格迁移的损失函数。 它由内容损失、风格损失和全变分损失3部分组成。

内容损失

与线性回归中的损失函数类似,内容损失通过平方误差函数衡量合成图像与内容图像在内容特征上的差异。 平方误差函数的两个输入均为extract_features函数计算所得到的内容层的输出。

def content_loss(Y_hat, Y):
    # 我们从动态计算梯度的树中分离目标:
    # 这是一个规定的值,而不是一个变量。
    return torch.square(Y_hat - Y.detach()).mean()

风格损失

风格损失与内容损失类似,也通过平方误差函数衡量合成图像与风格图像在风格上的差异。了表达风格层输出的风格,我们先通过extract_features函数计算风格层的输出。 假设该输出的样本数为 1 1 1,通道数为 c c c,高和宽分别为 h h h w w w,我们可以将此输出转换为矩阵 X X X,其有 c c c行和 w h wh wh列。 这个矩阵可以被看作由 c c c个长度为 w h wh wh的向量 x 1 , . . . , x c x_1,...,x_c x1,...,xc组合而成的。其中向量 x i x_i xi代表了通道 i i i上的风格特征。

在这些向量的格拉姆矩阵 X X T ∈ R c × c XX^T \in R^{c\times c} XXTRc×c中, i i i j j j列的元素 x i j x_{ij} xij即向量 x i x_i xi x j x_j xj的内积。它表达了通道 i i i和通道 j j j上风格特征的相关性。我们用这样的格拉姆矩阵来表达风格层输出的风格。 需要注意的是,当 h w hw hw的值较大时,格拉姆矩阵中的元素容易出现较大的值。 此外,格拉姆矩阵的高和宽皆为通道数 c c c。 为了让风格损失不受这些值的大小影响,下面定义的gram函数将格拉姆矩阵除以了矩阵中元素的个数,即 c h w chw chw

def gram(X):
    num_channels, n = X.shape[1], X.numel() // X.shape[1]
    X = X.reshape((num_channels, n))
    return torch.matmul(X, X.T) / (num_channels * n)

这段代码定义了一个函数 gram,它接受一个张量 X 作为输入,并返回 X 的 Gram 矩阵。Gram 矩阵是一种用于描述特征之间相关性的矩阵,通常用于计算输入图像的风格特征。

具体来说,这个函数首先获取输入张量 X 的通道数和像素数量,并将其重塑为一个形状为 (num_channels, n) 的矩阵,其中 num_channels 是通道数,n 是像素数量。然后,函数计算矩阵乘积 X @ X.T,并将结果除以 (num_channels * n),以归一化 Gram 矩阵的值。最后,函数返回归一化的 Gram 矩阵。

自然地,风格损失的平方误差函数的两个格拉姆矩阵输入分别基于合成图像与风格图像的风格层输出。这里假设基于风格图像的格拉姆矩阵gram_Y已经预先计算好了

def style_loss(Y_hat, gram_Y):
    return torch.square(gram(Y_hat) - gram_Y.detach()).mean()

全变分损失

有时候,我们学到的合成图像里面有大量高频噪点,即有特别亮或者特别暗的颗粒像素。 一种常见的去噪方法是全变分去噪(total variation denoising): 假设 x i , j x_{i,j} xi,j表示坐标 ( i , j ) (i,j) (i,j)处的像素值,降低全变分损失
∑ i , j ∣ x i , j − x i + 1 , j ∣ + ∣ x i , j − x i , j + 1 ∣ \sum_{i,j} |x_{i,j}-x_{i+1,j}|+|x_{i,j}-x_{i,j+1}| i,jxi,jxi+1,j+xi,jxi,j+1
能够尽可能使邻近的像素值相似。

def tv_loss(Y_hat):
    return 0.5 * (torch.abs(Y_hat[:, :, 1:, :] - Y_hat[:, :, :-1, :]).mean() +
                  torch.abs(Y_hat[:, :, :, 1:] - Y_hat[:, :, :, :-1]).mean())

总损失函数

风格转移的损失函数是内容损失、风格损失和总变化损失的加权和。 通过调节这些权重超参数,我们可以权衡合成图像在保留内容、迁移风格以及去噪三方面的相对重要性。

content_weight, style_weight, tv_weight = 1, 1e3, 10

def compute_loss(X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram):
    # 分别计算内容损失、风格损失和全变分损失
    contents_l = [content_loss(Y_hat, Y) * content_weight for Y_hat, Y in zip(
        contents_Y_hat, contents_Y)]
    styles_l = [style_loss(Y_hat, Y) * style_weight for Y_hat, Y in zip(
        styles_Y_hat, styles_Y_gram)]
    tv_l = tv_loss(X) * tv_weight
    # 对所有损失求和
    l = sum(10 * styles_l + contents_l + [tv_l])
    return contents_l, styles_l, tv_l, l

初始化合成图像

在风格迁移中,合成的图像是训练期间唯一需要更新的变量。因此,我们可以定义一个简单的模型SynthesizedImage,并将合成的图像视为模型参数。模型的前向传播只需返回模型参数即可。

class SynthesizedImage(nn.Module):
    def __init__(self, img_shape, **kwargs):
        super(SynthesizedImage, self).__init__(**kwargs)
        self.weight = nn.Parameter(torch.rand(*img_shape))

    def forward(self):
        return self.weight

下面,我们定义get_inits函数。该函数创建了合成图像的模型实例,并将其初始化为图像X。风格图像在各个风格层的格拉姆矩阵styles_Y_gram将在训练前预先计算好。

def get_inits(X, device, lr, styles_Y):
    gen_img = SynthesizedImage(X.shape).to(device)
    gen_img.weight.data.copy_(X.data)
    trainer = torch.optim.Adam(gen_img.parameters(), lr=lr)
    styles_Y_gram = [gram(Y) for Y in styles_Y]
    return gen_img(), styles_Y_gram, trainer

训练模型

def train(X, contents_Y, styles_Y, device, lr, num_epochs, lr_decay_epoch):
    X, styles_Y_gram, trainer = get_inits(X, device, lr, styles_Y)
    scheduler = torch.optim.lr_scheduler.StepLR(trainer, lr_decay_epoch, 0.8)
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[10, num_epochs],
                            legend=['content', 'style', 'TV'],
                            ncols=2, figsize=(7, 2.5))
    for epoch in range(num_epochs):
        trainer.zero_grad()
        contents_Y_hat, styles_Y_hat = extract_features(
            X, content_layers, style_layers)
        contents_l, styles_l, tv_l, l = compute_loss(
            X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram)
        l.backward()
        trainer.step()
        scheduler.step()
        if (epoch + 1) % 10 == 0:
            animator.axes[1].imshow(postprocess(X))
            animator.add(epoch + 1, [float(sum(contents_l)),
                                     float(sum(styles_l)), float(tv_l)])
    return X

现在我们训练模型: 首先将内容图像和风格图像的高和宽分别调整为300和450像素,用内容图像来初始化合成图像。

device, image_shape = d2l.try_gpu(), (300, 450)
net = net.to(device)
content_X, contents_Y = get_contents(image_shape, device)
_, styles_Y = get_styles(image_shape, device)
output = train(content_X, contents_Y, styles_Y, device, 0.3, 500, 50)

在这里插入图片描述

猜你喜欢

转载自blog.csdn.net/qq_51957239/article/details/131085254