Python如何实现图像风格化？神经风格迁移

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

# 1. 加载图像
content_image = Image.open("content.jpg").convert('RGB')
style_image = Image.open("style.jpg").convert('RGB')

# 2. 预处理
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])

content_image = transform(content_image).unsqueeze(0) # 添加batch维度
style_image = transform(style_image).unsqueeze(0)

# 3. 加载预训练的VGG19模型
model = models.vgg19(pretrained=True).features.eval() # 只使用特征提取部分
for param in model.parameters():
    param.requires_grad = False  # 冻结参数

# 4. 定义特征提取函数 (简化版)
def get_features(image, model):
    features = {}
    x = image
    for name, layer in model._modules.items():
        x = layer(x)
        if name in ['4', '9', '18', '27', '36']: # 示例层，实际需要根据VGG19结构调整
            features[name] = x
    return features

content_features = get_features(content_image, model)
style_features = get_features(style_image, model)

# 5. 定义损失函数 (简化版)
def content_loss(gen_features, content_features):
    return torch.mean((gen_features - content_features)**2)

def style_loss(gen_features, style_features):
    batch_size, channel, height, width = gen_features.size()
    G = torch.mm(gen_features.view(channel, height * width),
                   gen_features.view(channel, height * width).t()) # Gram矩阵
    A = torch.mm(style_features.view(channel, height * width),
                   style_features.view(channel, height * width).t())
    return torch.mean((G - A)**2) / (4 * channel**2 * height**2 * width**2)

# 6. 优化 (简化版)
generated_image = content_image.clone().requires_grad_(True) # 从内容图像开始
optimizer = optim.Adam([generated_image], lr=0.003)

for i in range(100): # 迭代次数
    gen_features = get_features(generated_image, model)

    content_loss_val = content_loss(gen_features['18'], content_features['18']) # 示例层
    style_loss_val = style_loss(gen_features['4'], style_features['4']) + \
                     style_loss(gen_features['9'], style_features['9']) + \
                     style_loss(gen_features['18'], style_features['18']) + \
                     style_loss(gen_features['27'], style_features['27']) + \
                     style_loss(gen_features['36'], style_features['36'])

    total_loss = content_loss_val * 1 + style_loss_val * 1e4 # 权重需要调整

    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()

    print(f"Iteration {i}, Loss: {total_loss.item()}")

# 7. 后处理 (简化版)
generated_image = generated_image.clone().detach().squeeze()
generated_image = generated_image.to('cpu').numpy()
generated_image = generated_image.transpose((1, 2, 0))
generated_image = generated_image * np.array((0.229, 0.224, 0.225)) + np.array((0.485, 0.456, 0.406))
generated_image = generated_image.clip(0, 1)

import matplotlib.pyplot as plt
plt.imshow(generated_image)
plt.show()