| import os | |
| from tqdm import tqdm | |
| import torch | |
| from torch.utils.data import DataLoader | |
| from logger import Logger, Visualizer | |
| import numpy as np | |
| import imageio | |
| from sync_batchnorm import DataParallelWithCallback | |
| def reconstruction(config, generator, kp_detector, checkpoint, log_dir, dataset): | |
| png_dir = os.path.join(log_dir, 'reconstruction/png') | |
| log_dir = os.path.join(log_dir, 'reconstruction') | |
| if checkpoint is not None: | |
| Logger.load_cpk(checkpoint, generator=generator, kp_detector=kp_detector) | |
| else: | |
| raise AttributeError("Checkpoint should be specified for mode='reconstruction'.") | |
| dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1) | |
| if not os.path.exists(log_dir): | |
| os.makedirs(log_dir) | |
| if not os.path.exists(png_dir): | |
| os.makedirs(png_dir) | |
| loss_list = [] | |
| if torch.cuda.is_available(): | |
| generator = DataParallelWithCallback(generator) | |
| kp_detector = DataParallelWithCallback(kp_detector) | |
| generator.eval() | |
| kp_detector.eval() | |
| for it, x in tqdm(enumerate(dataloader)): | |
| if config['reconstruction_params']['num_videos'] is not None: | |
| if it > config['reconstruction_params']['num_videos']: | |
| break | |
| with torch.no_grad(): | |
| predictions = [] | |
| visualizations = [] | |
| if torch.cuda.is_available(): | |
| x['video'] = x['video'].cuda() | |
| kp_source = kp_detector(x['video'][:, :, 0]) | |
| for frame_idx in range(x['video'].shape[2]): | |
| source = x['video'][:, :, 0] | |
| driving = x['video'][:, :, frame_idx] | |
| kp_driving = kp_detector(driving) | |
| out = generator(source, kp_source=kp_source, kp_driving=kp_driving) | |
| out['kp_source'] = kp_source | |
| out['kp_driving'] = kp_driving | |
| del out['sparse_deformed'] | |
| predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0]) | |
| visualization = Visualizer(**config['visualizer_params']).visualize(source=source, | |
| driving=driving, out=out) | |
| visualizations.append(visualization) | |
| loss_list.append(torch.abs(out['prediction'] - driving).mean().cpu().numpy()) | |
| predictions = np.concatenate(predictions, axis=1) | |
| imageio.imsave(os.path.join(png_dir, x['name'][0] + '.png'), (255 * predictions).astype(np.uint8)) | |
| image_name = x['name'][0] + config['reconstruction_params']['format'] | |
| imageio.mimsave(os.path.join(log_dir, image_name), visualizations) | |
| print("Reconstruction loss: %s" % np.mean(loss_list)) | |