escapist413/StyleFusion

1"""2风格迁移算法 - StyleFusion3作者: hammershock4描述: 该脚本用于实现非实时风格迁移，通过深度学习模型融合内容图像和风格图像。5      本实现依赖于PyTorch框架，并使用预训练的VGG模型作为风格和内容特征的提取器。6 7使用方法:8    运行该脚本时，可以通过命令行参数指定内容图像路径、风格图像路径、输出目录等。9    示例: python stylefusion.py --content_image ./data/city.jpg --style_image ./data/udnie.jpg --output_dir ./output/iterative_style_transfer10"""11 12import argparse13import torch14import torch.optim as optim15from tqdm import tqdm16 17from models import VGG18from utils import make_transform, load_image, save_image, calculate_style_loss, calculate_content_loss19 20 21def parse_args():22    parser = argparse.ArgumentParser(description='风格迁移算法 - StyleFusion')23    parser.add_argument('--content_image', type=str, default='./data/city.jpg', help='内容图像路径')24    parser.add_argument('--style_image', type=str, default='./data/udnie.jpg', help='风格图像路径')25    parser.add_argument('--output_dir', type=str, default='./output/iterative_style_transfer', help='输出目录')26    parser.add_argument('--image_size', type=int, nargs=2, default=[300, 450], help='图像尺寸 (高度, 宽度)')27    parser.add_argument('--content_weight', type=float, default=1, help='内容权重')28    parser.add_argument('--style_weight', type=float, default=15, help='风格权重')29    parser.add_argument('--epochs', type=int, default=20, help='训练周期数')30    parser.add_argument('--steps_per_epoch', type=int, default=100, help='每个周期的步数')31    parser.add_argument('--learning_rate', type=float, default=0.03, help='学习率')32    return parser.parse_args()33 34 35if __name__ == "__main__":36    args = parse_args()37    # 内容特征层及loss加权系数38    content_layers = {'5': 0.5, '10': 0.5}39    # 风格特征层及loss加权系数40    style_layers = {'0': 0.2, '5': 0.2, '10': 0.2, '19': 0.2, '28': 0.2}41 42    # ----------------训练即推理过程----------------43    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")44    transform = make_transform(args.image_size, normalize=True)45 46    content_img = load_image(args.content_image, transform).to(device)  # 读取内容图像47    style_img = load_image(args.style_image, transform).to(device)  # 读取风格图像48 49    generated_img = torch.randn_like(content_img, requires_grad=True).to(device)  # 随机初始化目标图像50    save_image(generated_img, args.output_dir, 'noise.jpg')  # 保存初始噪声图，供查看51 52    vgg_model = VGG(content_layers, style_layers).to(device).eval()  # 实例化模型和优化器53    # print(vgg_model.model)  # 打印vgg模型结构54    optimizer = optim.Adam([generated_img], lr=args.learning_rate)  # 直接对图像本身进行优化55 56    content_features, _ = vgg_model(content_img)  # 计算内容图的内容特征57    _, style_features = vgg_model(style_img)  # 计算风格图的风格特征58 59    for epoch in range(args.epochs):60        p_bar = tqdm(range(args.steps_per_epoch), desc=f'epoch {epoch}')61        for step in p_bar:62            generated_content, generated_style = vgg_model(generated_img)  # 计算生成图片的不同层次的内容特征和风格特征63            # 不同层次的内容和风格特征损失加权求和64            content_loss = sum(args.content_weight * content_layers[name] * calculate_content_loss(content_features[name], gen_content) for name, gen_content in generated_content.items())65            style_loss = sum(args.style_weight * style_layers[name] * calculate_style_loss(style_features[name], gen_style) for name, gen_style in generated_style.items())66 67            total_loss = style_loss + content_loss68 69            optimizer.zero_grad()70            total_loss.backward()71            torch.nn.utils.clip_grad_norm_(generated_img, max_norm=1)  # 梯度裁剪72            optimizer.step()73            p_bar.set_postfix(style_loss=style_loss.item(), content_loss=content_loss.item())74 75        # 保存生成图像76        save_image(generated_img, args.output_dir, f'generated_{epoch + 1}.jpg')77 78