# Load experiment setting
with open(opts.config, 'r') as stream:
config = yaml.load(stream, Loader=yaml.FullLoader)
# Activate GPUs
config['gpu_ids'] = opts.gpu_ids
gpu_info = trainer_util.activate_gpus(config)
# Get data loaders
cfg_val_loader = config['val_dataloader']
cfg_test_loader1 = config['test_dataloader1']
cfg_test_loader2 = config['test_dataloader2']
cfg_test_loader3 = config['test_dataloader3']
cfg_test_loader4 = config['test_dataloader4']
val_loader = trainer_util.get_dataloader(cfg_val_loader['dataset_args'],
cfg_val_loader['dataloader_args'])
test_loader1 = trainer_util.get_dataloader(cfg_test_loader1['dataset_args'],
cfg_test_loader1['dataloader_args'])
test_loader2 = trainer_util.get_dataloader(cfg_test_loader2['dataset_args'],
cfg_test_loader2['dataloader_args'])
test_loader3 = trainer_util.get_dataloader(cfg_test_loader3['dataset_args'],
cfg_test_loader3['dataloader_args'])
test_loader4 = trainer_util.get_dataloader(cfg_test_loader4['dataset_args'],
cfg_test_loader4['dataloader_args'])
# get model
if 'vgg' in config['model']['architecture'] and 'guided' in config['model'][
'architecture']:
diss_model = GuidedDissimNet(**config['model']).cuda()
if 'vgg' in config['model']['architecture'] and 'correlated' in config[
'model']['architecture']:
def mae_features(config_file_path, gpu_ids, dataroot, data_origin):
soft_fdr = os.path.join(dataroot, 'mae_features_' + data_origin)
if not os.path.exists(soft_fdr):
os.makedirs(soft_fdr)
# load experiment setting
with open(config_file_path, 'r') as stream:
config = yaml.load(stream, Loader=yaml.FullLoader)
# activate GPUs
config['gpu_ids'] = gpu_ids
gpu = int(gpu_ids)
# get data_loaders
cfg_test_loader = config['test_dataloader']
cfg_test_loader['dataset_args']['dataroot'] = dataroot
test_loader = trainer_util.get_dataloader(cfg_test_loader['dataset_args'], cfg_test_loader['dataloader_args'])
class VGG19(torch.nn.Module):
def __init__(self, requires_grad=False):
super().__init__()
vgg_pretrained_features = torchvision.models.vgg19(pretrained=True).features
self.slice1 = torch.nn.Sequential()
self.slice2 = torch.nn.Sequential()
self.slice3 = torch.nn.Sequential()
self.slice4 = torch.nn.Sequential()
self.slice5 = torch.nn.Sequential()
for x in range(2):
self.slice1.add_module(str(x), vgg_pretrained_features[x])
for x in range(2, 7):
self.slice2.add_module(str(x), vgg_pretrained_features[x])
for x in range(7, 12):
self.slice3.add_module(str(x), vgg_pretrained_features[x])
for x in range(12, 21):
self.slice4.add_module(str(x), vgg_pretrained_features[x])
for x in range(21, 30):
self.slice5.add_module(str(x), vgg_pretrained_features[x])
if not requires_grad:
for param in self.parameters():
param.requires_grad = False
def forward(self, X):
h_relu1 = self.slice1(X)
h_relu2 = self.slice2(h_relu1)
h_relu3 = self.slice3(h_relu2)
h_relu4 = self.slice4(h_relu3)
h_relu5 = self.slice5(h_relu4)
out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
return out
from torch.nn.modules.upsampling import Upsample
up5 = Upsample(scale_factor=16, mode='bicubic')
up4 = Upsample(scale_factor=8, mode='bicubic')
up3 = Upsample(scale_factor=4, mode='bicubic')
up2 = Upsample(scale_factor=2, mode='bicubic')
up1 = Upsample(scale_factor=1, mode='bicubic')
to_pil = ToPILImage()
# Going through visualization loader
weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
vgg = VGG19().cuda(gpu)
with torch.no_grad():
for i, data_i in enumerate(test_loader):
print('Generating image %i out of %i'%(i+1, len(test_loader)))
img_name = os.path.basename(data_i['original_path'][0])
original = data_i['original'].cuda(gpu)
synthesis = data_i['synthesis'].cuda(gpu)
x_vgg, y_vgg = vgg(original), vgg(synthesis)
feat5 = torch.mean(torch.abs(x_vgg[4] - y_vgg[4]), dim=1).unsqueeze(1)
feat4 = torch.mean(torch.abs(x_vgg[3] - y_vgg[3]), dim=1).unsqueeze(1)
feat3 = torch.mean(torch.abs(x_vgg[2] - y_vgg[2]), dim=1).unsqueeze(1)
feat2 = torch.mean(torch.abs(x_vgg[1] - y_vgg[1]), dim=1).unsqueeze(1)
feat1 = torch.mean(torch.abs(x_vgg[0] - y_vgg[0]), dim=1).unsqueeze(1)
img_5 = up5(feat5)
img_4 = up4(feat4)
img_3 = up3(feat3)
img_2 = up2(feat2)
img_1 = up1(feat1)
combined = weights[0] * img_1 + weights[1] * img_2 + weights[2] * img_3 + weights[3] * img_4 + weights[
4] * img_5
min_v = torch.min(combined.squeeze())
max_v = torch.max(combined.squeeze())
combined = (combined.squeeze() - min_v) / (max_v - min_v)
combined = to_pil(combined.cpu())
pred_name = 'mea_' + img_name
combined.save(os.path.join(soft_fdr, pred_name))
if not os.path.isdir(os.path.join(store_fdr_exp, 'pred')):
os.makedirs(os.path.join(store_fdr_exp, 'label'), exist_ok=True)
os.makedirs(os.path.join(store_fdr_exp, 'pred'), exist_ok=True)
os.makedirs(os.path.join(store_fdr_exp, 'soft'), exist_ok=True)
# Activate GPUs
config['gpu_ids'] = opts.gpu_ids
gpu_info = trainer_util.activate_gpus(config)
# checks if we are using prior images
prior = config['model']['prior']
# Get data loaders
cfg_test_loader = config['test_dataloader']
# adds logic to dataloaders (avoid repetition in config file)
cfg_test_loader['dataset_args']['prior'] = prior
test_loader = trainer_util.get_dataloader(cfg_test_loader['dataset_args'],
cfg_test_loader['dataloader_args'])
# get model
if config['model']['prior']:
diss_model = DissimNetPrior(**config['model']).cuda()
elif 'vgg' in config['model']['architecture']:
diss_model = DissimNet(**config['model']).cuda()
else:
raise NotImplementedError()
diss_model.eval()
model_path = os.path.join(save_fdr, exp_name,
'%s_net_%s.pth' % (epoch, exp_name))
model_weights = torch.load(model_path)
diss_model.load_state_dict(model_weights)