用python写一个图片风格迁移的代码

好的，这里给出一个使用PyTorch实现的图片风格迁移的代码。

首先需要安装必要的库：`torch`, `torchvision`, `numpy`, `PIL`

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

# 定义参数
style_img_path = 'style.jpg'  # 风格图像路径
content_img_path = 'content.jpg'  # 内容图像路径
output_img_path = 'output.jpg'  # 输出图像路径
image_size = 512  # 图像大小
style_weight = 100000  # 风格权重
content_weight = 1  # 内容权重
num_epochs = 2000  # 迭代次数

# 加载图像并进行预处理
def load_image(image_path, transform=None, max_size=None, shape=None):
    image = Image.open(image_path)
    if max_size is not None:
        scale = max_size / max(image.size)
        size = np.array(image.size) * scale
        size = tuple(map(int, size))
        image = image.resize(size, Image.ANTIALIAS)
    if shape is not None:
        image = image.resize(shape, Image.LANCZOS)
    if transform is not None:
        image = transform(image).unsqueeze(0)
    return image

# 将图像转换为张量，并进行归一化
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])
style_img = load_image(style_img_path, transform, max_size=image_size)
content_img = load_image(content_img_path, transform, shape=[style_img.shape[2], style_img.shape[3]])
output_img = content_img.clone().requires_grad_(True)

# 加载预训练的VGG模型
vgg = models.vgg19(pretrained=True).features

# 将模型转移到GPU上（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vgg.to(device)
style_img.to(device)
content_img.to(device)
output_img.to(device)

# 定义风格损失函数
class StyleLoss(nn.Module):
    def __init__(self, target_feature):
        super(StyleLoss, self).__init__()
        self.target = self.gram_matrix(target_feature).detach()

    def gram_matrix(self, 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)

    def forward(self, input):
        G = self.gram_matrix(input)
        self.loss = nn.functional.mse_loss(G, self.target)
        return input

# 定义内容损失函数
class ContentLoss(nn.Module):
    def __init__(self, target_feature):
        super(ContentLoss, self).__init__()
        self.target = target_feature.detach()

    def forward(self, input):
        self.loss = nn.functional.mse_loss(input, self.target)
        return input

# 提取特征图函数
def get_features(image, model):
    layers = {
        '0': 'conv1_1',
        '5': 'conv2_1',
        '10': 'conv3_1',
        '19': 'conv4_1',
        '21': 'conv4_2',
        '28': 'conv5_1'
    }
    features = {}
    x = image
    for name, layer in model._modules.items():
        x = layer(x)
        if name in layers:
            features[layers[name]] = x
    return features

# 优化器
optimizer = optim.Adam([output_img], lr=0.01)

# 迭代训练
for epoch in range(num_epochs):
    output_features = get_features(output_img, vgg)
    content_features = get_features(content_img, vgg)
    style_features = get_features(style_img, vgg)

    style_loss = 0
    content_loss = 0

    # 计算风格损失
    for layer in style_features:
        target_feature = style_features[layer]
        target_gram = StyleLoss(target_feature)
        output_feature = output_features[layer]
        output_gram = StyleLoss(output_feature)
        style_loss += target_gram.loss * style_weight

    # 计算内容损失
    for layer in content_features:
        target_feature = content_features[layer]
        output_feature = output_features[layer]
        content_loss += ContentLoss(target_feature)(output_feature) * content_weight

    # 总损失
    total_loss = style_loss + content_loss

    # 更新输出图像
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()

    # 输出日志
    if (epoch + 1) % 100 == 0:
        print('Epoch [{}/{}], Style Loss: {:.4f}, Content Loss: {:.4f}, Total Loss: {:.4f}'
              .format(epoch + 1, num_epochs, style_loss.item(), content_loss.item(), total_loss.item()))

# 保存输出图像
output_img = output_img.cpu().detach().squeeze(0)
output_img = output_img.permute(1, 2, 0)
output_img = output_img * torch.tensor([0.229, 0.224, 0.225]) + torch.tensor([0.485, 0.456, 0.406])
output_img = output_img.numpy()
output_img = np.clip(output_img, 0, 1)
output_img = (output_img * 255).astype(np.uint8)
output_img = Image.fromarray(output_img)
output_img.save(output_img_path)