python风格迁移_图片风格迁移(Neural style)

图片风格迁移训练时使用预处理的模型，不进行模型的训练，将图片像素值设置为可训练参数，进行图片的训练。

结合一张图片的内容和另一张图片的风格，生成一张新图片

论文地址：https://arxiv.org/pdf/1508.06576.pdf

损失函数：

以下代码通过将图片x输入到预训练好的vgg19模型中，取模型features模块的0, 5, 10, 19, 28五层的输出值计算contentLoss与StyleLoss之和，对该loss进行反向传播，从而修改图片的像素值。

"""

图片风格迁移，使用预训练的vgg，不训练vgg模型，而是进行图片的训练，训练我们想要的风格图片

"""

import torch

from PIL import Image

import argparse

import torch.nn as nn

import matplotlib.pyplot as plt

import numpy as np

import torchvision

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_image(image_path, transform=None, max_size=None, shape=None):

image = Image.open(image_path)

if max_size:

scale = max_size/max(image.size)

size = np.array(image.size)*scale

image = image.resize(size.astype(int), Image.ANTIALIAS)

if shape:

image = image.resize(shape, Image.LANCZOS)

if transform:

image = transform(image).unsqueeze(0)

return image.to(device)

transform = torchvision.transforms.Compose([

torchvision.transforms.ToTensor(),

# mean and std based on trainDatasets

torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

])

content = load_image("data/content.jpg", transform, max_size=400)

style = load_image("data/style.jpg", transform, shape=[content.size(2), content.size(3)])

print(content.shape) # torch.Size([1, 3, 400, 300])

print(style.shape) # torch.Size([1, 3, 300, 400])

def imshow(tensor, title=None):

image = tensor.cpu().clone()

image = image.squeeze(0)

# ToPILImage为类对象，因此需要ToPILImage()(image)

image = torchvision.transforms.ToPILImage()(image)

plt.imshow(image)

if title is not None:

plt.title(title)

plt.pause(2) # pause a bit so that plots are updated

class VGGNet(nn.Module):

def __init__(self):

super(VGGNet, self).__init__()

self.select = ['0', '5', '10', '19', '28']

self.vgg = torchvision.models.vgg19(pretrained=True).features

def forward(self, x):

features = []

for name, layer in self.vgg._modules.items():

x = layer(x)

if name in self.select:

features.append(x)

return features

# 不训练模型，eval()模式

vgg = VGGNet().to(device).eval()

for feat in vgg(content):

print(feat.shape)

# 训练图片

target = content.clone().requires_grad_(True)

print("target shape: ", target.shape)

optimizer = torch.optim.Adam([target], lr=0.003, betas=[0.5, 0.999])

num_steps = 2000

for step in range(num_steps):

target_features = vgg(target)

content_features = vgg(content)

style_features = vgg(style)

content_loss = style_loss = 0.

for f1, f2, f3 in zip(target_features, content_features, style_features):

content_loss += torch.mean((f1-f2)**2)

_, c, h, w = f1.size()

f1 = f1.view(c, h*w)

f3 = f3.view(c, h*w)

f1 = torch.mm(f1, f1.t())

f3 = torch.mm(f3, f3.t())

style_loss += torch.mean((f1-f3)**2)/(c*h*w)

loss = content_loss + style_loss*100.

optimizer.zero_grad()

loss.backward()

optimizer.step()

if step % 100 == 0:

print("Step [{}/{}, Content Loss: {:.4f}, Style Loss: {:.4f}"

.format(step, num_steps, content_loss.item(), style_loss.item()))

denorm = torchvision.transforms.Normalize([-2.12, -2.04, -1.8], [4.37, 4.46, 4.44])

img = target.clone().squeeze()

img = denorm(img).clamp_(0, 1)

imshow(img, title="Target Image")

plt.pause(100)