pytorch 图像风格迁移

PyTorch是一个流行的深度学习框架，可以用于图像风格迁移。图像风格迁移是将一张图像的风格应用到另一张图像上的过程。以下是实现图像风格迁移的一些步骤：

1. 准备数据集：准备一组内容图像和一组风格图像。

2. 定义损失函数：定义内容损失和风格损失，用于衡量生成图像与内容图像和风格图像之间的差异。

3. 定义模型：定义一个卷积神经网络模型，用于将内容图像转换为风格图像。

4. 训练模型：使用数据集训练模型，以最小化损失函数。

5. 进行风格迁移：使用训练好的模型将内容图像转换为风格图像。

以下是一个简单的PyTorch图像风格迁移的例子：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image

# 加载图像
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)

# 定义损失函数
class ContentLoss(nn.Module):
    def __init__(self, target):
        super(ContentLoss, self).__init__()
        self.target = target.detach()
    def forward(self, input):
        self.loss = F.mse_loss(input, self.target)
        return input

class StyleLoss(nn.Module):
    def __init__(self, target_feature):
        super(StyleLoss, self).__init__()
        self.target = gram_matrix(target_feature).detach()
    def forward(self, input):
        G = gram_matrix(input)
        self.loss = F.mse_loss(G, self.target)
        return input

def gram_matrix(input):
    a, b, c, d = input.size()
    features = input.view(a * b, c * d)
    G = torch.mm(features, features.t())
    return G.div(a * b * c * d)

# 定义模型
class TransformerNet(nn.Module):
    def __init__(self):
        super(TransformerNet, self).__init__()
        self.conv1 = ConvLayer(3, 32, kernel_size=9, stride=1)
        self.in1 = nn.InstanceNorm2d(32, affine=True)
        self.conv2 = ConvLayer(32, 64, kernel_size=3, stride=2)
        self.in2 = nn.InstanceNorm2d(64, affine=True)
        self.conv3 = ConvLayer(64, 128, kernel_size=3, stride=2)
        self.in3 = nn.InstanceNorm2d(128, affine=True)
        self.res1 = ResidualBlock(128)
        self.res2 = ResidualBlock(128)
        self.res3 = ResidualBlock(128)
        self.res4 = ResidualBlock(128)
        self.res5 = ResidualBlock(128)
        self.conv4 = ConvLayer(128, 64, kernel_size=3, stride=1)
        self.in4 = nn.InstanceNorm2d(64, affine=True)
        self.conv5 = ConvLayer(64, 3, kernel_size=9, stride=1)
    def forward(self, input):
        x = F.relu(self.in1(self.conv1(input)))
        x = F.relu(self.in2(self.conv2(x)))
        x = F.relu(self.in3(self.conv3(x)))
        x = self.res1(x)
        x = self.res2(x)
        x = self.res3(x)
        x = self.res4(x)
        x = self.res5(x)
        x = F.relu(self.in4(self.conv4(x)))
        x = self.conv5(x)
        return x

class ConvLayer(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride):
        super(ConvLayer, self).__init__()
        reflection_padding = kernel_size // 2
        self.reflection_pad = nn.ReflectionPad2d(reflection_padding)
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)
    def forward(self, x):
        out = self.reflection_pad(x)
        out = self.conv2d(out)
        return out

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1)
        self.in1 = nn.InstanceNorm2d(channels, affine=True)
        self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1)
        self.in2 = nn.InstanceNorm2d(channels, affine=True)
        self.relu = nn.ReLU()
    def forward(self, x):
        residual = x
        out = self.relu(self.in1(self.conv1(x)))
        out = self.in2(self.conv2(out))
        out = out + residual
        return out

# 训练模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(256),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])])
content = load_image("content.jpg", transform, max_size=400)
style = load_image("style.jpg", transform, shape=[content.size(2), content.size(3)])
model = TransformerNet().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
style_features = models.vgg19(pretrained=True).features(style.to(device))
style_grams = [gram_matrix(feature) for feature in style_features]
for i in range(1, 501):
    optimizer.zero_grad()
    content_features = model(content)
    content_loss = 0
    style_loss = 0
    for feature in content_features:
        content_loss += ContentLoss(feature).loss
    for j, feature in enumerate(model(style)):
        style_loss += StyleLoss(feature).loss
    total_loss = content_loss + style_loss
    total_loss.backward()
    optimizer.step()
    if i % 50 == 0:
        print("Iteration:", i, "Total loss:", total_loss.item())

# 进行风格迁移
with torch.no_grad():
    output = model(content)
output_image = output.cpu().clone()
output_image = output_image.squeeze(0)
output_image = transforms.ToPILImage()(output_image)
output_image.save("output.jpg")