def main(args):
if args.device is None:
device = torch.device('cpu')
else:
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device)
device = torch.device('cuda:0')
torch.cuda.set_device(0)
# Load XP config
xp_config = load_json(os.path.join(args.xp_dir, 'params.json'))
xp_config.device = device
xp_config.data_dir = args.data_dir
xp_config.xp_dir = args.xp_dir
xp_config.nt_pred = 10
args.nt_pred = 10
test_set = load_dataset(xp_config, train=False)
sep_net = load_model(xp_config, args.epoch)
all_mse, all_ssim = compute_mse_ssim(xp_config, test_set, sep_net)
mse_array = np.concatenate(all_mse, axis=0)
ssim_array = np.concatenate(all_ssim, axis=0)
print(f'MSE at t+10: {np.mean(mse_array.mean(axis=0)[:10])}')
print(f'MSE at t+6: {np.mean(mse_array.mean(axis=0)[:6])}')
print(f'SSIM at t+10: {np.mean(ssim_array.mean(axis=0)[:10])}')
print(f'SSIM at t+6: {np.mean(ssim_array.mean(axis=0)[:6])}')
def main(args):
if args.device is None:
device = torch.device('cpu')
else:
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device)
device = torch.device('cuda:0')
torch.cuda.set_device(0)
# Load XP config
xp_config = load_json(os.path.join(args.xp_dir, 'params.json'))
xp_config.device = device
xp_config.nt_pred = 40
args.nt_pred = 40
test_set = load_dataset(xp_config, train=False)
sep_net = build_model(xp_config)
all_mse = compute_mse(xp_config, args.batch_size, test_set, sep_net)
mse_array = np.concatenate(all_mse, axis=0)
print(f'MSE at t+40: {np.mean(mse_array.mean(axis=0)[:40])}')
def main(args):
##################################################################################################################
# Setup
##################################################################################################################
# -- Device handling (CPU, GPU)
if args.device is None:
device = torch.device('cpu')
else:
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device)
device = torch.device('cuda:0')
torch.cuda.set_device(0)
# Seed
random.seed(args.test_seed)
np.random.seed(args.test_seed)
torch.manual_seed(args.test_seed)
# Load XP config
xp_config = load_json(os.path.join(args.xp_dir, 'params.json'))
xp_config.device = device
xp_config.data_dir = args.data_dir
xp_config.xp_dir = args.xp_dir
xp_config.nt_pred = args.nt_pred
##################################################################################################################
# Load test data
##################################################################################################################
print('Loading data...')
test_dataset = load_dataset(xp_config, train=False)
test_loader = DataLoader(test_dataset,
batch_size=args.batch_size,
pin_memory=True)
train_dataset = load_dataset(xp_config, train=True)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
pin_memory=True)
nc = 1
size = 64
##################################################################################################################
# Load model
##################################################################################################################
print('Loading model...')
sep_net = build_model(xp_config)
##################################################################################################################
# Eval
##################################################################################################################
print('Generating samples...')
torch.set_grad_enabled(False)
train_iterator = iter(train_loader)
nt_test = xp_config.nt_cond + args.nt_pred
predictions = []
content_swap = []
cond_swap = []
target_swap = []
cond = []
gt = []
results = defaultdict(list)
# Evaluation is done by batch
for batch in tqdm(test_loader, ncols=80, desc='evaluation'):
# Data
x_cond, x_target = batch
bsz = len(x_cond)
x_cond = x_cond.to(device)
x_target = x_target.to(device)
cond.append(x_cond.cpu().mul(255).byte().permute(0, 1, 3, 4, 2))
gt.append(x_target.cpu().mul(255).byte().permute(0, 1, 3, 4, 2))
# Prediction
x_pred, _, s_codes, _ = sep_net.get_forecast(x_cond, nt_test)
x_pred = x_pred[:, xp_config.nt_cond:]
# Content swap
x_swap_cond, x_swap_target = next(train_iterator)
x_swap_cond = x_swap_cond[:bsz].to(device)
x_swap_target = x_swap_target[:bsz].to(device)
x_swap_cond_byte = x_swap_cond.cpu().mul(255).byte()
x_swap_target_byte = x_swap_target.cpu().mul(255).byte()
cond_swap.append(x_swap_cond_byte.permute(0, 1, 3, 4, 2))
target_swap.append(x_swap_target_byte.permute(0, 1, 3, 4, 2))
x_swap_pred = sep_net.get_forecast(x_swap_cond,
nt_test,
init_s_code=s_codes[:, 0])[0]
x_swap_pred = x_swap_pred[:, xp_config.dt:]
content_swap.append(x_swap_pred.cpu().mul(255).byte().permute(
0, 1, 3, 4, 2))
# Pixelwise quantitative eval
x_target = x_target.view(-1, args.nt_pred, nc, size, size)
mse = torch.mean(F.mse_loss(x_pred, x_target, reduction='none'),
dim=[3, 4])
metrics_batch = {
'mse': mse.mean(2).mean(1).cpu(),
'psnr': 10 * torch.log10(1 / mse).mean(2).mean(1).cpu(),
'ssim': _ssim_wrapper(x_pred, x_target).mean(2).mean(1).cpu()
}
predictions.append(x_pred.cpu().mul(255).byte().permute(0, 1, 3, 4, 2))
# Compute metrics for best samples and register
for name in metrics_batch.keys():
results[name].append(metrics_batch[name])
##################################################################################################################
# Print results
##################################################################################################################
print('\n')
print('Results:')
for name in results.keys():
res = torch.cat(results[name]).numpy()
results[name] = res
print(name, res.mean(), '+/-', 1.960 * res.std() / np.sqrt(len(res)))
##################################################################################################################
# Save samples
##################################################################################################################
np.savez_compressed(os.path.join(args.xp_dir, 'results.npz'), **results)
np.savez_compressed(os.path.join(args.xp_dir, 'predictions.npz'),
predictions=torch.cat(predictions).numpy())
np.savez_compressed(os.path.join(args.xp_dir, 'gt.npz'),
gt=torch.cat(gt).numpy())
np.savez_compressed(os.path.join(args.xp_dir, 'cond.npz'),
cond=torch.cat(cond).numpy())
np.savez_compressed(os.path.join(args.xp_dir, 'content_swap.npz'),
content_swap=torch.cat(content_swap).numpy())
np.savez_compressed(os.path.join(args.xp_dir, 'cond_swap.npz'),
target_swap=torch.cat(cond_swap).numpy())
np.savez_compressed(os.path.join(args.xp_dir, 'target_swap.npz'),
target_swap=torch.cat(target_swap).numpy())
def main(args):
##################################################################################################################
# Setup
##################################################################################################################
# -- Device handling (CPU, GPU)
if args.device is None:
device = torch.device('cpu')
else:
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device)
device = torch.device('cuda:0')
torch.cuda.set_device(0)
# Seed
random.seed(args.test_seed)
np.random.seed(args.test_seed)
torch.manual_seed(args.test_seed)
# Load XP config
xp_config = load_json(os.path.join(args.xp_dir, 'params.json'))
xp_config.device = device
xp_config.data_dir = args.data_dir
xp_config.xp_dir = args.xp_dir
xp_config.nt_pred = args.nt_pred
xp_config.n_object = 1
##################################################################################################################
# Load test data
##################################################################################################################
print('Loading data...')
test_dataset = load_dataset(xp_config, train=False)
test_loader = DataLoader(test_dataset,
batch_size=args.batch_size,
pin_memory=True)
swap_dataset = SwapDataset(False,
args.data_dir,
xp_config.nt_cond,
seq_len=xp_config.nt_cond + args.nt_pred)
swap_loader = DataLoader(swap_dataset,
batch_size=args.batch_size,
pin_memory=True)
nc = 3
size = 64
##################################################################################################################
# Load model
##################################################################################################################
print('Loading model...')
sep_net = load_model(xp_config, args.epoch)
##################################################################################################################
# Eval
##################################################################################################################
print('Generating samples...')
torch.set_grad_enabled(False)
swap_iterator = iter(swap_loader)
nt_test = xp_config.nt_cond + args.nt_pred
gt_swap = []
content_swap = []
cond_swap = []
target_swap = []
results = defaultdict(list)
# Evaluation is done by batch
for batch in tqdm(test_loader, ncols=80, desc='evaluation'):
# Data
x_cond, x_target, _, x_gt_swap = next(swap_iterator)
x_gt_swap = x_gt_swap.to(device)
x_cond = x_cond.to(device)
# Extraction of S
_, _, s_code, _ = sep_net.get_forecast(x_cond, nt_test)
# Content swap
x_swap_cond, x_swap_target = batch
x_swap_cond = x_swap_cond.to(device)
x_swap_target = x_swap_target.to(device)
x_swap_cond_byte = x_cond.cpu().mul(255).byte()
x_swap_target_byte = x_swap_target.cpu().mul(255).byte()
cond_swap.append(x_swap_cond_byte.permute(0, 1, 3, 4, 2))
target_swap.append(x_swap_target_byte.permute(0, 1, 3, 4, 2))
x_swap_pred = sep_net.get_forecast(x_swap_cond,
nt_test,
init_s_code=s_code)[0]
x_swap_pred = x_swap_pred[:, xp_config.nt_cond:]
content_swap.append(x_swap_pred.cpu().mul(255).byte().permute(
0, 1, 3, 4, 2))
gt_swap.append(x_gt_swap[:, 0].cpu().mul(255).byte().permute(
0, 1, 3, 4, 2))
# Pixelwise quantitative eval
x_gt_swap = x_gt_swap.view(-1, xp_config.n_object, args.nt_pred, nc,
size, size).to(device)
metrics_batch = {'mse': [], 'psnr': [], 'ssim': []}
for j, reordering in enumerate(
itertools.permutations(range(xp_config.n_object))):
mse = torch.mean(F.mse_loss(x_swap_pred,
x_gt_swap[:, j],
reduction='none'),
dim=[3, 4])
metrics_batch['mse'].append(mse.mean(2).mean(1).cpu())
metrics_batch['psnr'].append(
10 * torch.log10(1 / mse).mean(2).mean(1).cpu())
metrics_batch['ssim'].append(
_ssim_wrapper(x_swap_pred, x_gt_swap[:,
j]).mean(2).mean(1).cpu())
# Compute metrics for best samples and register
results['mse'].append(
torch.min(torch.stack(metrics_batch['mse']), 0)[0])
results['psnr'].append(
torch.max(torch.stack(metrics_batch['psnr']), 0)[0])
results['ssim'].append(
torch.max(torch.stack(metrics_batch['ssim']), 0)[0])
##################################################################################################################
# Print results
##################################################################################################################
print('\n')
print('Results:')
for name in results.keys():
res = torch.cat(results[name]).numpy()
results[name] = res
print(name, res.mean())
##################################################################################################################
# Save samples
##################################################################################################################
np.savez_compressed(os.path.join(args.xp_dir, 'results_swap.npz'),
**results)
np.savez_compressed(os.path.join(args.xp_dir, 'content_swap_gt.npz'),
gt_swap=torch.cat(gt_swap).numpy())
np.savez_compressed(os.path.join(args.xp_dir, 'content_swap_test.npz'),
content_swap=torch.cat(content_swap).numpy())
np.savez_compressed(os.path.join(args.xp_dir, 'cond_swap_test.npz'),
cond_swap=torch.cat(cond_swap).numpy())
np.savez_compressed(os.path.join(args.xp_dir, 'target_swap_test.npz'),
target_swap=torch.cat(target_swap).numpy())