dipfa/StyleTransfer-main/comparison.py

# -*- coding: utf-8 -*- 
""" 
@File : comparison.py 
@Author: csc
@Date : 2022/6/28
"""
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '5'

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import torchvision
import torchvision.transforms as transforms
import torchvision.models as models

from tensorboardX import SummaryWriter

import random
import shutil
from glob import glob
from tqdm import tqdm

from utils import *
from models import *

device = 'cpu'


class Config:
    def __init__(self, name, backbone_name, transform_net, meta_path, content_path, style_path, transform_path, style_weight):
        self.name = name
        self.backbone_name = backbone_name
        if backbone_name == 'vgg16':
            vgg16 = models.vgg16(pretrained=False)
            vgg16.load_state_dict(torch.load('./models/vgg16-397923af.pth'))
            self.backbone = VGG(vgg16.features[:23]).to(device).eval()
        elif backbone_name == 'vgg19':
            vgg19 = models.vgg19(pretrained=False)
            vgg19.load_state_dict(torch.load('./models/vgg19-dcbb9e9d.pth'))
            self.backbone = VGG19(vgg19.features[:30]).to(device).eval()
        self.transform_net = transform_net
        self.meta_path = meta_path
        self.transform_path = transform_path
        self.content_path = content_path
        self.style_path = style_path
        self.style_weight = style_weight
        self.transformed_images = None


base = 32
configs = [
    # Config(name='vgg16_500',
    #        backbone_name='vgg16',
    #        transform_net=TransformNet(base).to(device),
    #        meta_path='../weight/new_500_22ep/metanet_base32_style50_tv1e-06_tagnohvd.pth',
    #        transform_path='../weight/new_500_22ep/metanet_base32_style50_tv1e-06_tagnohvd_transform_net.pth',
    #        content_path='../COCO2014_500/',
    #        style_path='../WikiArt_500/',
    #        style_weight=50),
    # Config(name='vgg16_1000',
    #        backbone_name='vgg16',
    #        transform_net=TransformNet(base).to(device),
    #        meta_path='../weight/1000_22ep/metanet_base32_style50_tv1e-06_tagnohvd.pth',
    #        transform_path='../weight/1000_22ep/metanet_base32_style50_tv1e-06_tagnohvd_transform_net.pth',
    #        content_path='../COCO2014_1000/',
    #        style_path='../WikiArt_1000/',
    #        style_weight=50),
    # Config(name='resnext_gram_500',
    #        backbone_name='vgg16',
    #        transform_net=TransformNet(base, residuals='resnext').to(device),
    #        meta_path='../weight/resnext_gram_500_22ep_3e5/metanet_base32_style300000.0_tv1e-06_tagnohvd.pth',
    #        transform_path='../weight/resnext_gram_500_22ep_3e5/metanet_base32_style300000.0_tv1e-06_tagnohvd_transform_net.pth',
    #        content_path='../COCO2014_500/',
    #        style_path='../WikiArt_500/',
    #        style_weight=3e5),
    # Config(name='resnext_gram_1000',
    #        backbone_name='vgg16',
    #        transform_net=TransformNet(base, residuals='resnext').to(device),
    #        meta_path='../weight/resnext_gram_1000_22ep_3e5/metanet_base32_style300000.0_tv1e-06_tagnohvd.pth',
    #        transform_path='../weight/resnext_gram_1000_22ep_3e5/metanet_base32_style300000.0_tv1e-06_tagnohvd_transform_net.pth',
    #        content_path='../COCO2014_1000/',
    #        style_path='../WikiArt_1000/',
    #        style_weight=3e5),
    # Config(name='vgg19_resnext_gram_500',
    #        backbone_name='vgg19',
    #        transform_net=TransformNet(base, residuals='resnext').to(device),
    #        meta_path='../weight/vgg19_resnext_gram_500/metanet_base32_style300000.0_tv1e-06_tagnohvd.pth',
    #        transform_path='../weight/vgg19_resnext_gram_500/metanet_base32_style300000.0_tv1e-06_tagnohvd_transform_net.pth',
    #        content_path='../COCO2014_500/',
    #        style_path='../WikiArt_500/',
    #        style_weight=3e5),
    Config(name='vgg19_resnext_gram_1000',
           backbone_name='vgg19',
           transform_net=TransformNet(base, residuals='resnext').to(device),
           meta_path='../weight/vgg19_resnext_gram_1000/metanet_base32_style300000.0_tv1e-06_tagnohvd.pth',
           transform_path='../weight/vgg19_resnext_gram_1000/metanet_base32_style300000.0_tv1e-06_tagnohvd_transform_net.pth',
           content_path='../COCO2014_1000/',
           style_path='../WikiArt_1000/',
           style_weight=3e5),
    # 消融
    # Config(name='vgg19_resnet_gram_500',
    #        backbone_name='vgg19',
    #        transform_net=TransformNet(base).to(device),
    #        meta_path='../weight/vgg19_resnet_gram_500/metanet_base32_style300000.0_tv1e-06_tagnohvd.pth',
    #        transform_path='../weight/vgg19_resnet_gram_500/metanet_base32_style300000.0_tv1e-06_tagnohvd_transform_net.pth',
    #        content_path='../COCO2014_500/',
    #        style_path='../WikiArt_500/',
    #        style_weight=3e5),
    # Config(name='vgg19_resnext_mse_500',
    #        backbone_name='vgg19',
    #        transform_net=TransformNet(base, residuals='resnext').to(device),
    #        meta_path='../weight/vgg19_resnext_mse_500/metanet_base32_style50_tv1e-06_tagnohvd.pth',
    #        transform_path='../weight/vgg19_resnext_mse_500/metanet_base32_style50_tv1e-06_tagnohvd_transform_net.pth',
    #        content_path='../COCO2014_500/',
    #        style_path='../WikiArt_500/',
    #        style_weight=50),
]
content_weight = 1
tv_weight = 1e-6
batch_size = 8

# 可视化
width = 256
data_transform = transforms.Compose([
    transforms.RandomResizedCrop(width, scale=(256/480, 1), ratio=(1, 1)),
    transforms.ToTensor(),
    tensor_normalizer
])


class MetaNet(nn.Module):
    def __init__(self, param_dict, backbone='vgg16'):
        super(MetaNet, self).__init__()
        self.param_num = len(param_dict)
        if backbone == 'vgg16':
            self.hidden = nn.Linear(1920, 128 * self.param_num)
        elif backbone == 'vgg19':
            self.hidden = nn.Linear(1920, 128 * self.param_num)
        self.fc_dict = {}
        for i, (name, params) in enumerate(param_dict.items()):
            self.fc_dict[name] = i
            setattr(self, 'fc{}'.format(i + 1), nn.Linear(128, params))

    # ONNX 要求输出 tensor 或者 list，不能是 dict
    def forward(self, mean_std_features):
        hidden = F.relu(self.hidden(mean_std_features))
        filters = {}
        for name, i in self.fc_dict.items():
            fc = getattr(self, 'fc{}'.format(i + 1))
            filters[name] = fc(hidden[:, i * 128:(i + 1) * 128])
        return list(filters.values())

    def forward2(self, mean_std_features):
        hidden = F.relu(self.hidden(mean_std_features))
        filters = {}
        for name, i in self.fc_dict.items():
            fc = getattr(self, 'fc{}'.format(i + 1))
            filters[name] = fc(hidden[:, i * 128:(i + 1) * 128])
        return filters

graph = [[] for i in range(3)]

for (index, config) in enumerate(configs):
    backbone_name = config.backbone_name
    backbone = config.backbone
    transform_net = config.transform_net
    meta_path = config.meta_path
    transform_path = config.transform_path
    content_path = config.content_path
    style_path = config.style_path
    style_weight = config.style_weight

    style_dataset = torchvision.datasets.ImageFolder(style_path, transform=data_transform)
    content_dataset = torchvision.datasets.ImageFolder(content_path, transform=data_transform)

    content_data_loader = torch.utils.data.DataLoader(content_dataset, batch_size=batch_size, shuffle=True)

    # style_img_name = '30925'
    # test_style_image = read_image(style_path + 'train/' + style_img_name + '.jpg', target_width=width).to(device)
    test_style_image = read_image('./images/pearl.jpg', target_width=width).to(device)

    style_features = backbone(test_style_image)
    style_mean_std = mean_std(style_features)
    metanet = MetaNet(transform_net.get_param_dict(), backbone=backbone_name).to(device)

    trainable_params = {}
    trainable_param_shapes = {}
    for model in [backbone, transform_net, metanet]:
        for name, param in model.named_parameters():
            if param.requires_grad:
                trainable_params[name] = param
                trainable_param_shapes[name] = param.shape

    optimizer = optim.Adam(trainable_params.values(), 1e-3)

    metanet.load_state_dict(torch.load(meta_path))
    transform_net.load_state_dict(torch.load(transform_path))

    n_batch = 20
    with tqdm(enumerate(content_data_loader), total=n_batch) as pbar:
        for batch, (content_images, _) in pbar:
            x = content_images.cpu().numpy()
            if (x.min(-1).min(-1) == x.max(-1).max(-1)).any():
                continue

            optimizer.zero_grad()

            # 使用风格图像生成风格模型
            weights = metanet.forward2(mean_std(style_features))
            transform_net.set_weights(weights, 0)

            # 使用风格模型预测风格迁移图像
            content_images = content_images.to(device)
            transformed_images = transform_net(content_images)

            # 计算特征
            content_features = backbone(content_images)
            transformed_features = backbone(transformed_images)
            transformed_mean_std = mean_std(transformed_features)

            # content loss
            content_loss = content_weight * F.mse_loss(transformed_features[2], content_features[2])

            # style loss
            style_loss = style_weight * F.mse_loss(transformed_mean_std,
                                                   style_mean_std.expand_as(transformed_mean_std))

            # total variation loss
            y = transformed_images
            tv_loss = tv_weight * (torch.sum(torch.abs(y[:, :, :, :-1] - y[:, :, :, 1:])) +
                                   torch.sum(torch.abs(y[:, :, :-1, :] - y[:, :, 1:, :])))

            # 求和
            loss = content_loss + style_loss + tv_loss

            loss.backward()
            optimizer.step()

            if batch > n_batch:
                break

    content_img_path = content_path + 'train2014/'
    content_img_names = [
        content_img_path + 'COCO_train2014_000000081611.jpg',
        content_img_path + 'COCO_train2014_000000149739.jpg',
        content_img_path + 'COCO_train2014_000000505057.jpg',
        content_img_path + 'COCO_train2014_000000421773.jpg',
        './images/dancing.png',
        './images/boat.png',
        './images/ecnu.jpg',
        './images/text.jpg'
    ]
    test_content_images = torch.stack([read_image(name, target_width=width) for name in content_img_names]).to(device)
    content_images_vis = torch.cat([x for x in test_content_images], dim=-1)

    # config.transformed_images = transform_net(test_content_images)
    config.transformed_images = [transform_net(read_image(name, target_width=width).to(device)) for name in content_img_names]


fig = plt.figure(figsize=(50, 20))
size = 310
line_len = len(configs)
plt.subplot(size + 1)
imshow(test_style_image)
plt.subplot(size + line_len + 1)
imshow(content_images_vis)
for (index, config) in enumerate(configs):
    plt.subplot(size + line_len * (index + 2) + 1)
    plt.title(config.name)
    transformed_images_vis = torch.cat([x for x in config.transformed_images], dim=-1)
    imshow(transformed_images_vis)

fig.savefig('./images/out.png')