用python写一个图片风格化迁移项目

好的，以下是一个简单的图片风格化迁移项目示例，使用 Python 和 PyTorch 框架。

步骤一：导入必要的库

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import os

步骤二：定义超参数和数据加载器

# 定义超参数
content_weight = 1  # 内容损失函数的权重
style_weight = 1000  # 风格损失函数的权重
epochs = 20  # 迭代次数
batch_size = 4  # 批处理大小

# 加载数据集
data_transform = transforms.Compose([
    transforms.Resize(256),  # 缩放到256x256像素大小
    transforms.CenterCrop(256),  # 中心裁剪为256x256像素
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化
])
content_dataset = datasets.ImageFolder('./content', transform=data_transform)
style_dataset = datasets.ImageFolder('./style', transform=data_transform)
content_loader = DataLoader(content_dataset, batch_size=batch_size, shuffle=True)
style_loader = DataLoader(style_dataset, batch_size=batch_size, shuffle=True)

步骤三：定义模型

# 定义风格迁移模型
class StyleTransferModel(nn.Module):
    def __init__(self):
        super(StyleTransferModel, self).__init__()
        self.features = models.vgg19(pretrained=True).features[:35]  # 加载预训练的VGG19模型
        for param in self.parameters():
            param.requires_grad = False  # 冻结参数
        self.content_loss = nn.MSELoss()  # 定义内容损失函数
        self.style_loss = nn.MSELoss()  # 定义风格损失函数
        self.content_feature = None  # 内容图像的特征
        self.style_features = None  # 风格图像的特征
        self.target_feature = None  # 目标图像的特征

    def forward(self, x):
        self.content_feature = self.features(x.clone())  # 克隆一份x，防止直接修改导致误差计算错误
        return x

    def compute_content_loss(self):
        loss = self.content_loss(self.target_feature, self.content_feature)
        return content_weight * loss

    def compute_style_loss(self):
        loss = 0
        for i in range(len(self.style_features)):
            target_gram = self.gram_matrix(self.target_feature[i])
            style_gram = self.gram_matrix(self.style_features[i])
            loss += self.style_loss(target_gram, style_gram)
        return style_weight * loss

    def gram_matrix(self, x):
        b, c, h, w = x.size()
        features = x.view(b * c, h * w)
        G = torch.mm(features, features.t())
        return G.div(b * c * h * w)

    def set_style_features(self, x):
        self.style_features = []
        for feature in self.features:
            x = feature(x)
            if isinstance(feature, nn.ReLU):
                feature.inplace = False
            if isinstance(feature, nn.MaxPool2d):
                self.style_features.append(x)
                if len(self.style_features) == 5:
                    return

    def set_target_feature(self, x):
        self.target_feature = self.features(x.clone())

步骤四：定义训练函数

def train(model, content_loader, style_loader, epochs):
    optimizer = optim.Adam(model.parameters(), lr=0.001)  # 定义优化器
    for epoch in range(epochs):
        model.train()
        content_iter = iter(content_loader)
        style_iter = iter(style_loader)
        for i in range(len(content_iter)):
            content, _ = content_iter.next()
            style, _ = style_iter.next()
            model.set_style_features(style)  # 设置风格图像的特征
            model.set_target_feature(content)  # 设置目标图像的特征
            optimizer.zero_grad()  # 梯度清零
            loss = model.compute_content_loss() + model.compute_style_loss()  # 计算损失函数
            loss.backward()  # 反向传播
            optimizer.step()  # 更新参数
        print("Epoch ", epoch + 1, " complete.")

步骤五：定义测试函数

def test(model, content_path, style_path, output_path):
    content_image = Image.open(content_path)
    style_image = Image.open(style_path)
    transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    content = transform(content_image).unsqueeze(0)
    style = transform(style_image).unsqueeze(0)
    model.set_style_features(style)
    model.set_target_feature(content)
    output = model(content)
    output_image = output.squeeze().detach().numpy()
    output_image = np.transpose(output_image, (1, 2, 0))
    output_image = output_image * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]
    output_image = np.clip(output_image, 0, 1)
    output_image = Image.fromarray((output_image * 255).astype(np.uint8))
    output_image.save(output_path)

步骤六：训练模型

model = StyleTransferModel()
train(model, content_loader, style_loader, epochs)

步骤七：测试模型

test(model, './test_content.jpg', './test_style.jpg', './output.jpg')

以上是一个简单的图片风格化迁移项目示例，你可以根据需要进行修改或优化。