gulabpatel/First-Order-Motion

1import matplotlib2matplotlib.use('Agg')3import os, sys4import yaml5from argparse import ArgumentParser6from tqdm import tqdm7 8import imageio9import numpy as np10from skimage.transform import resize11from skimage import img_as_ubyte12import torch13from sync_batchnorm import DataParallelWithCallback14 15from modules.generator import OcclusionAwareGenerator16from modules.keypoint_detector import KPDetector17from animate import normalize_kp18from scipy.spatial import ConvexHull19 20 21if sys.version_info[0] < 3:22    raise Exception("You must use Python 3 or higher. Recommended version is Python 3.7")23 24def load_checkpoints(config_path, checkpoint_path, cpu=False):25 26    with open(config_path) as f:27        config = yaml.load(f)28 29    generator = OcclusionAwareGenerator(**config['model_params']['generator_params'],30                                        **config['model_params']['common_params'])31    if not cpu:32        generator.cuda()33 34    kp_detector = KPDetector(**config['model_params']['kp_detector_params'],35                             **config['model_params']['common_params'])36    if not cpu:37        kp_detector.cuda()38    39    if cpu:40        checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))41    else:42        checkpoint = torch.load(checkpoint_path)43 44    generator.load_state_dict(checkpoint['generator'])45    kp_detector.load_state_dict(checkpoint['kp_detector'])46    47    if not cpu:48        generator = DataParallelWithCallback(generator)49        kp_detector = DataParallelWithCallback(kp_detector)50 51    generator.eval()52    kp_detector.eval()53    54    return generator, kp_detector55 56 57def make_animation(source_image, driving_video, generator, kp_detector, relative=True, adapt_movement_scale=True, cpu=False):58    with torch.no_grad():59        predictions = []60        source = torch.tensor(source_image[np.newaxis].astype(np.float32)).permute(0, 3, 1, 2)61        if not cpu:62            source = source.cuda()63        driving = torch.tensor(np.array(driving_video)[np.newaxis].astype(np.float32)).permute(0, 4, 1, 2, 3)64        kp_source = kp_detector(source)65        kp_driving_initial = kp_detector(driving[:, :, 0])66 67        for frame_idx in tqdm(range(driving.shape[2])):68            driving_frame = driving[:, :, frame_idx]69            if not cpu:70                driving_frame = driving_frame.cuda()71            kp_driving = kp_detector(driving_frame)72            kp_norm = normalize_kp(kp_source=kp_source, kp_driving=kp_driving,73                                   kp_driving_initial=kp_driving_initial, use_relative_movement=relative,74                                   use_relative_jacobian=relative, adapt_movement_scale=adapt_movement_scale)75            out = generator(source, kp_source=kp_source, kp_driving=kp_norm)76 77            predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])78    return predictions79 80def find_best_frame(source, driving, cpu=False):81    import face_alignment82 83    def normalize_kp(kp):84        kp = kp - kp.mean(axis=0, keepdims=True)85        area = ConvexHull(kp[:, :2]).volume86        area = np.sqrt(area)87        kp[:, :2] = kp[:, :2] / area88        return kp89 90    fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=True,91                                      device='cpu' if cpu else 'cuda')92    kp_source = fa.get_landmarks(255 * source)[0]93    kp_source = normalize_kp(kp_source)94    norm  = float('inf')95    frame_num = 096    for i, image in tqdm(enumerate(driving)):97        kp_driving = fa.get_landmarks(255 * image)[0]98        kp_driving = normalize_kp(kp_driving)99        new_norm = (np.abs(kp_source - kp_driving) ** 2).sum()100        if new_norm < norm:101            norm = new_norm102            frame_num = i103    return frame_num104 105if __name__ == "__main__":106    parser = ArgumentParser()107    parser.add_argument("--config", required=True, help="path to config")108    parser.add_argument("--checkpoint", default='vox-cpk.pth.tar', help="path to checkpoint to restore")109 110    parser.add_argument("--source_image", default='sup-mat/source.png', help="path to source image")111    parser.add_argument("--driving_video", default='sup-mat/source.png', help="path to driving video")112    parser.add_argument("--result_video", default='result.mp4', help="path to output")113 114    parser.add_argument("--relative", dest="relative", action="store_true", help="use relative or absolute keypoint coordinates")115    parser.add_argument("--adapt_scale", dest="adapt_scale", action="store_true", help="adapt movement scale based on convex hull of keypoints")116 117    parser.add_argument("--find_best_frame", dest="find_best_frame", action="store_true", 118                        help="Generate from the frame that is the most alligned with source. (Only for faces, requires face_aligment lib)")119 120    parser.add_argument("--best_frame", dest="best_frame", type=int, default=None,  121                        help="Set frame to start from.")122 123    parser.add_argument("--cpu", dest="cpu", action="store_true", help="cpu mode.")124 125 126    parser.set_defaults(relative=False)127    parser.set_defaults(adapt_scale=False)128 129    opt = parser.parse_args()130 131    source_image = imageio.imread(opt.source_image)132    reader = imageio.get_reader(opt.driving_video)133    fps = reader.get_meta_data()['fps']134    driving_video = []135    try:136        for im in reader:137            driving_video.append(im)138    except RuntimeError:139        pass140    reader.close()141 142    source_image = resize(source_image, (256, 256))[..., :3]143    driving_video = [resize(frame, (256, 256))[..., :3] for frame in driving_video]144    generator, kp_detector = load_checkpoints(config_path=opt.config, checkpoint_path=opt.checkpoint, cpu=opt.cpu)145 146    if opt.find_best_frame or opt.best_frame is not None:147        i = opt.best_frame if opt.best_frame is not None else find_best_frame(source_image, driving_video, cpu=opt.cpu)148        print ("Best frame: " + str(i))149        driving_forward = driving_video[i:]150        driving_backward = driving_video[:(i+1)][::-1]151        predictions_forward = make_animation(source_image, driving_forward, generator, kp_detector, relative=opt.relative, adapt_movement_scale=opt.adapt_scale, cpu=opt.cpu)152        predictions_backward = make_animation(source_image, driving_backward, generator, kp_detector, relative=opt.relative, adapt_movement_scale=opt.adapt_scale, cpu=opt.cpu)153        predictions = predictions_backward[::-1] + predictions_forward[1:]154    else:155        predictions = make_animation(source_image, driving_video, generator, kp_detector, relative=opt.relative, adapt_movement_scale=opt.adapt_scale, cpu=opt.cpu)156    imageio.mimsave(opt.result_video, [img_as_ubyte(frame) for frame in predictions], fps=fps)157 158