|
|
# -*- 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')
|
