| import os |
| from tqdm import tqdm |
|
|
| import torch |
| from torch.utils.data import DataLoader |
|
|
| from frames_dataset import PairedDataset |
| from logger import Logger, Visualizer |
| import imageio |
| from scipy.spatial import ConvexHull |
| import numpy as np |
|
|
| from sync_batchnorm import DataParallelWithCallback |
|
|
|
|
| def normalize_kp(kp_source, kp_driving, kp_driving_initial, adapt_movement_scale=False, |
| use_relative_movement=False, use_relative_jacobian=False): |
| if adapt_movement_scale: |
| source_area = ConvexHull(kp_source['value'][0].data.cpu().numpy()).volume |
| driving_area = ConvexHull(kp_driving_initial['value'][0].data.cpu().numpy()).volume |
| adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area) |
| else: |
| adapt_movement_scale = 1 |
|
|
| kp_new = {k: v for k, v in kp_driving.items()} |
|
|
| if use_relative_movement: |
| kp_value_diff = (kp_driving['value'] - kp_driving_initial['value']) |
| kp_value_diff *= adapt_movement_scale |
| kp_new['value'] = kp_value_diff + kp_source['value'] |
|
|
| if use_relative_jacobian: |
| jacobian_diff = torch.matmul(kp_driving['jacobian'], torch.inverse(kp_driving_initial['jacobian'])) |
| kp_new['jacobian'] = torch.matmul(jacobian_diff, kp_source['jacobian']) |
|
|
| return kp_new |
|
|
|
|
| def animate(config, generator, kp_detector, checkpoint, log_dir, dataset): |
| log_dir = os.path.join(log_dir, 'animation') |
| png_dir = os.path.join(log_dir, 'png') |
| animate_params = config['animate_params'] |
|
|
| dataset = PairedDataset(initial_dataset=dataset, number_of_pairs=animate_params['num_pairs']) |
| dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1) |
|
|
| if checkpoint is not None: |
| Logger.load_cpk(checkpoint, generator=generator, kp_detector=kp_detector) |
| else: |
| raise AttributeError("Checkpoint should be specified for mode='animate'.") |
|
|
| if not os.path.exists(log_dir): |
| os.makedirs(log_dir) |
|
|
| if not os.path.exists(png_dir): |
| os.makedirs(png_dir) |
|
|
| 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)): |
| with torch.no_grad(): |
| predictions = [] |
| visualizations = [] |
|
|
| driving_video = x['driving_video'] |
| source_frame = x['source_video'][:, :, 0, :, :] |
|
|
| kp_source = kp_detector(source_frame) |
| kp_driving_initial = kp_detector(driving_video[:, :, 0]) |
|
|
| for frame_idx in range(driving_video.shape[2]): |
| driving_frame = driving_video[:, :, frame_idx] |
| kp_driving = kp_detector(driving_frame) |
| kp_norm = normalize_kp(kp_source=kp_source, kp_driving=kp_driving, |
| kp_driving_initial=kp_driving_initial, **animate_params['normalization_params']) |
| out = generator(source_frame, kp_source=kp_source, kp_driving=kp_norm) |
|
|
| out['kp_driving'] = kp_driving |
| out['kp_source'] = kp_source |
| out['kp_norm'] = kp_norm |
|
|
| 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_frame, |
| driving=driving_frame, out=out) |
| visualization = visualization |
| visualizations.append(visualization) |
|
|
| predictions = np.concatenate(predictions, axis=1) |
| result_name = "-".join([x['driving_name'][0], x['source_name'][0]]) |
| imageio.imsave(os.path.join(png_dir, result_name + '.png'), (255 * predictions).astype(np.uint8)) |
|
|
| image_name = result_name + animate_params['format'] |
| imageio.mimsave(os.path.join(log_dir, image_name), visualizations) |
|
|
