Spaces:
Runtime error
Runtime error
File size: 4,964 Bytes
280b585 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
import matplotlib
matplotlib.use('Agg')
import sys
import yaml
from argparse import ArgumentParser
from tqdm import tqdm
from scipy.spatial import ConvexHull
import numpy as np
import imageio
from skimage.transform import resize
from skimage import img_as_ubyte
import torch
from modules.inpainting_network import InpaintingNetwork
from modules.keypoint_detector import KPDetector
from modules.dense_motion import DenseMotionNetwork
from modules.avd_network import AVDNetwork
def load_checkpoints(config_path, checkpoint_path, device):
with open(config_path) as f:
config = yaml.full_load(f)
inpainting = InpaintingNetwork(**config['model_params']['generator_params'],
**config['model_params']['common_params'])
kp_detector = KPDetector(**config['model_params']['common_params'])
dense_motion_network = DenseMotionNetwork(**config['model_params']['common_params'],
**config['model_params']['dense_motion_params'])
avd_network = AVDNetwork(num_tps=config['model_params']['common_params']['num_tps'],
**config['model_params']['avd_network_params'])
kp_detector.to(device)
dense_motion_network.to(device)
inpainting.to(device)
avd_network.to(device)
checkpoint = torch.load(checkpoint_path, map_location=device)
inpainting.load_state_dict(checkpoint['inpainting_network'])
kp_detector.load_state_dict(checkpoint['kp_detector'])
dense_motion_network.load_state_dict(checkpoint['dense_motion_network'])
if 'avd_network' in checkpoint:
avd_network.load_state_dict(checkpoint['avd_network'])
inpainting.eval()
kp_detector.eval()
dense_motion_network.eval()
avd_network.eval()
return inpainting, kp_detector, dense_motion_network, avd_network
def relative_kp(kp_source, kp_driving, kp_driving_initial):
source_area = ConvexHull(kp_source['fg_kp'][0].data.cpu().numpy()).volume
driving_area = ConvexHull(kp_driving_initial['fg_kp'][0].data.cpu().numpy()).volume
adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
kp_new = {k: v for k, v in kp_driving.items()}
kp_value_diff = (kp_driving['fg_kp'] - kp_driving_initial['fg_kp'])
kp_value_diff *= adapt_movement_scale
kp_new['fg_kp'] = kp_value_diff + kp_source['fg_kp']
return kp_new
def make_animation(source_image, driving_video, inpainting_network, kp_detector, dense_motion_network, avd_network, device, mode = 'relative'):
assert mode in ['standard', 'relative', 'avd']
with torch.no_grad():
predictions = []
source = torch.tensor(source_image[np.newaxis].astype(np.float32)).permute(0, 3, 1, 2)
source = source.to(device)
driving = torch.tensor(np.array(driving_video)[np.newaxis].astype(np.float32)).permute(0, 4, 1, 2, 3).to(device)
kp_source = kp_detector(source)
kp_driving_initial = kp_detector(driving[:, :, 0])
for frame_idx in tqdm(range(driving.shape[2])):
driving_frame = driving[:, :, frame_idx]
driving_frame = driving_frame.to(device)
kp_driving = kp_detector(driving_frame)
if mode == 'standard':
kp_norm = kp_driving
elif mode=='relative':
kp_norm = relative_kp(kp_source=kp_source, kp_driving=kp_driving,
kp_driving_initial=kp_driving_initial)
elif mode == 'avd':
kp_norm = avd_network(kp_source, kp_driving)
dense_motion = dense_motion_network(source_image=source, kp_driving=kp_norm,
kp_source=kp_source, bg_param = None,
dropout_flag = False)
out = inpainting_network(source, dense_motion)
predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])
return predictions
def find_best_frame(source, driving, cpu):
import face_alignment
def normalize_kp(kp):
kp = kp - kp.mean(axis=0, keepdims=True)
area = ConvexHull(kp[:, :2]).volume
area = np.sqrt(area)
kp[:, :2] = kp[:, :2] / area
return kp
fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=True,
device= 'cpu' if cpu else 'cuda')
kp_source = fa.get_landmarks(255 * source)[0]
kp_source = normalize_kp(kp_source)
norm = float('inf')
frame_num = 0
for i, image in tqdm(enumerate(driving)):
try:
kp_driving = fa.get_landmarks(255 * image)[0]
kp_driving = normalize_kp(kp_driving)
new_norm = (np.abs(kp_source - kp_driving) ** 2).sum()
if new_norm < norm:
norm = new_norm
frame_num = i
except:
pass
return frame_num |