feat: updates

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 yoyo-nb
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

app.py

@@ -0,0 +1,100 @@
+import torch
+import imageio
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+from skimage.transform import resize
+import warnings
+import os
+from demo import make_animation
+from skimage import img_as_ubyte
+from demo import load_checkpoints
+import gradio
+
+
+def inference(source_image_path='./assets/source.png', driving_video_path='./assets/driving.mp4', dataset_name="vox"):
+    # edit the config
+    device = torch.device('cpu')
+    # dataset_name = 'vox'  # ['vox', 'taichi', 'ted', 'mgif']
+    # source_image_path = './assets/source.png'
+    # driving_video_path = './assets/driving.mp4'
+    output_video_path = './generated.mp4'
+
+    pixel = 256  # for vox, taichi and mgif, the resolution is 256*256
+    if (dataset_name == 'ted'):  # for ted, the resolution is 384*384
+        pixel = 384
+    config_path = f'config/{dataset_name}-{pixel}.yaml'
+    checkpoint_path = f'checkpoints/{dataset_name}.pth.tar'
+    predict_mode = 'relative'  # ['standard', 'relative', 'avd']
+
+    warnings.filterwarnings("ignore")
+
+    source_image = imageio.imread(source_image_path)
+    reader = imageio.get_reader(driving_video_path)
+
+    source_image = resize(source_image, (pixel, pixel))[..., :3]
+
+    fps = reader.get_meta_data()['fps']
+    driving_video = []
+    try:
+        for im in reader:
+            driving_video.append(im)
+    except RuntimeError:
+        pass
+    reader.close()
+
+    driving_video = [resize(frame, (pixel, pixel))[..., :3] for frame in driving_video]
+
+    # driving_video = driving_video[:10]
+
+    def display(source, driving, generated=None) -> animation.ArtistAnimation:
+        fig = plt.figure(figsize=(8 + 4 * (generated is not None), 6))
+
+        ims = []
+        for i in range(len(driving)):
+            cols = [source]
+            cols.append(driving[i])
+            if generated is not None:
+                cols.append(generated[i])
+            im = plt.imshow(np.concatenate(cols, axis=1), animated=True)
+            plt.axis('off')
+            ims.append([im])
+
+        ani = animation.ArtistAnimation(fig, ims, interval=50, repeat_delay=1000)
+        # plt.show()
+        plt.close()
+        return ani
+
+    inpainting, kp_detector, dense_motion_network, avd_network = load_checkpoints(config_path=config_path,
+                                                                                  checkpoint_path=checkpoint_path,
+                                                                                  device=device)
+
+    predictions = make_animation(source_image, driving_video, inpainting, kp_detector, dense_motion_network,
+                                 avd_network, device=device, mode=predict_mode)
+
+    # save resulting video
+    imageio.mimsave(output_video_path, [img_as_ubyte(frame) for frame in predictions], fps=fps)
+
+    ani = display(source_image, driving_video, predictions)
+    ani.save('animation.mp4', writer='imagemagick', fps=60)
+    return 'animation.mp4'
+
+
+demo = gradio.Interface(
+    fn=inference,
+    inputs=[
+        gradio.inputs.Image(type="filepath", label="Input image"),
+        gradio.inputs.Video(label="Input video"),
+        gradio.inputs.Dropdown(['vox', 'taichi', 'ted', 'mgif'], type="value", default="vox", label="Model",
+                               optional=False),
+
+    ],
+    outputs=["video"],
+    examples=[
+        ['./assets/source.png', './assets/driving.mp4', "vox"],
+        ['./assets/source_ted.png', './assets/driving_ted.mp4', "ted"],
+    ],
+)
+
+if __name__ == "__main__":
+    demo.launch()

augmentation.py

@@ -0,0 +1,344 @@
+"""
+Code from https://github.com/hassony2/torch_videovision
+"""
+
+import numbers
+
+import random
+import numpy as np
+import PIL
+
+from skimage.transform import resize, rotate
+import torchvision
+
+import warnings
+
+from skimage import img_as_ubyte, img_as_float
+
+
+def crop_clip(clip, min_h, min_w, h, w):
+    if isinstance(clip[0], np.ndarray):
+        cropped = [img[min_h:min_h + h, min_w:min_w + w, :] for img in clip]
+
+    elif isinstance(clip[0], PIL.Image.Image):
+        cropped = [
+            img.crop((min_w, min_h, min_w + w, min_h + h)) for img in clip
+            ]
+    else:
+        raise TypeError('Expected numpy.ndarray or PIL.Image' +
+                        'but got list of {0}'.format(type(clip[0])))
+    return cropped
+
+
+def pad_clip(clip, h, w):
+    im_h, im_w = clip[0].shape[:2]
+    pad_h = (0, 0) if h < im_h else ((h - im_h) // 2, (h - im_h + 1) // 2)
+    pad_w = (0, 0) if w < im_w else ((w - im_w) // 2, (w - im_w + 1) // 2)
+
+    return np.pad(clip, ((0, 0), pad_h, pad_w, (0, 0)), mode='edge')
+
+
+def resize_clip(clip, size, interpolation='bilinear'):
+    if isinstance(clip[0], np.ndarray):
+        if isinstance(size, numbers.Number):
+            im_h, im_w, im_c = clip[0].shape
+            # Min spatial dim already matches minimal size
+            if (im_w <= im_h and im_w == size) or (im_h <= im_w
+                                                   and im_h == size):
+                return clip
+            new_h, new_w = get_resize_sizes(im_h, im_w, size)
+            size = (new_w, new_h)
+        else:
+            size = size[1], size[0]
+
+        scaled = [
+            resize(img, size, order=1 if interpolation == 'bilinear' else 0, preserve_range=True,
+                   mode='constant', anti_aliasing=True) for img in clip
+            ]
+    elif isinstance(clip[0], PIL.Image.Image):
+        if isinstance(size, numbers.Number):
+            im_w, im_h = clip[0].size
+            # Min spatial dim already matches minimal size
+            if (im_w <= im_h and im_w == size) or (im_h <= im_w
+                                                   and im_h == size):
+                return clip
+            new_h, new_w = get_resize_sizes(im_h, im_w, size)
+            size = (new_w, new_h)
+        else:
+            size = size[1], size[0]
+        if interpolation == 'bilinear':
+            pil_inter = PIL.Image.NEAREST
+        else:
+            pil_inter = PIL.Image.BILINEAR
+        scaled = [img.resize(size, pil_inter) for img in clip]
+    else:
+        raise TypeError('Expected numpy.ndarray or PIL.Image' +
+                        'but got list of {0}'.format(type(clip[0])))
+    return scaled
+
+
+def get_resize_sizes(im_h, im_w, size):
+    if im_w < im_h:
+        ow = size
+        oh = int(size * im_h / im_w)
+    else:
+        oh = size
+        ow = int(size * im_w / im_h)
+    return oh, ow
+
+
+class RandomFlip(object):
+    def __init__(self, time_flip=False, horizontal_flip=False):
+        self.time_flip = time_flip
+        self.horizontal_flip = horizontal_flip
+
+    def __call__(self, clip):
+        if random.random() < 0.5 and self.time_flip:
+            return clip[::-1]
+        if random.random() < 0.5 and self.horizontal_flip:
+            return [np.fliplr(img) for img in clip]
+
+        return clip
+
+
+class RandomResize(object):
+    """Resizes a list of (H x W x C) numpy.ndarray to the final size
+    The larger the original image is, the more times it takes to
+    interpolate
+    Args:
+    interpolation (str): Can be one of 'nearest', 'bilinear'
+    defaults to nearest
+    size (tuple): (widht, height)
+    """
+
+    def __init__(self, ratio=(3. / 4., 4. / 3.), interpolation='nearest'):
+        self.ratio = ratio
+        self.interpolation = interpolation
+
+    def __call__(self, clip):
+        scaling_factor = random.uniform(self.ratio[0], self.ratio[1])
+
+        if isinstance(clip[0], np.ndarray):
+            im_h, im_w, im_c = clip[0].shape
+        elif isinstance(clip[0], PIL.Image.Image):
+            im_w, im_h = clip[0].size
+
+        new_w = int(im_w * scaling_factor)
+        new_h = int(im_h * scaling_factor)
+        new_size = (new_w, new_h)
+        resized = resize_clip(
+            clip, new_size, interpolation=self.interpolation)
+
+        return resized
+
+
+class RandomCrop(object):
+    """Extract random crop at the same location for a list of videos
+    Args:
+    size (sequence or int): Desired output size for the
+    crop in format (h, w)
+    """
+
+    def __init__(self, size):
+        if isinstance(size, numbers.Number):
+            size = (size, size)
+
+        self.size = size
+
+    def __call__(self, clip):
+        """
+        Args:
+        img (PIL.Image or numpy.ndarray): List of videos to be cropped
+        in format (h, w, c) in numpy.ndarray
+        Returns:
+        PIL.Image or numpy.ndarray: Cropped list of videos
+        """
+        h, w = self.size
+        if isinstance(clip[0], np.ndarray):
+            im_h, im_w, im_c = clip[0].shape
+        elif isinstance(clip[0], PIL.Image.Image):
+            im_w, im_h = clip[0].size
+        else:
+            raise TypeError('Expected numpy.ndarray or PIL.Image' +
+                            'but got list of {0}'.format(type(clip[0])))
+
+        clip = pad_clip(clip, h, w)
+        im_h, im_w = clip.shape[1:3]
+        x1 = 0 if h == im_h else random.randint(0, im_w - w)
+        y1 = 0 if w == im_w else random.randint(0, im_h - h)
+        cropped = crop_clip(clip, y1, x1, h, w)
+
+        return cropped
+
+
+class RandomRotation(object):
+    """Rotate entire clip randomly by a random angle within
+    given bounds
+    Args:
+    degrees (sequence or int): Range of degrees to select from
+    If degrees is a number instead of sequence like (min, max),
+    the range of degrees, will be (-degrees, +degrees).
+    """
+
+    def __init__(self, degrees):
+        if isinstance(degrees, numbers.Number):
+            if degrees < 0:
+                raise ValueError('If degrees is a single number,'
+                                 'must be positive')
+            degrees = (-degrees, degrees)
+        else:
+            if len(degrees) != 2:
+                raise ValueError('If degrees is a sequence,'
+                                 'it must be of len 2.')
+
+        self.degrees = degrees
+
+    def __call__(self, clip):
+        """
+        Args:
+        img (PIL.Image or numpy.ndarray): List of videos to be cropped
+        in format (h, w, c) in numpy.ndarray
+        Returns:
+        PIL.Image or numpy.ndarray: Cropped list of videos
+        """
+        angle = random.uniform(self.degrees[0], self.degrees[1])
+        if isinstance(clip[0], np.ndarray):
+            rotated = [rotate(image=img, angle=angle, preserve_range=True) for img in clip]
+        elif isinstance(clip[0], PIL.Image.Image):
+            rotated = [img.rotate(angle) for img in clip]
+        else:
+            raise TypeError('Expected numpy.ndarray or PIL.Image' +
+                            'but got list of {0}'.format(type(clip[0])))
+
+        return rotated
+
+
+class ColorJitter(object):
+    """Randomly change the brightness, contrast and saturation and hue of the clip
+    Args:
+    brightness (float): How much to jitter brightness. brightness_factor
+    is chosen uniformly from [max(0, 1 - brightness), 1 + brightness].
+    contrast (float): How much to jitter contrast. contrast_factor
+    is chosen uniformly from [max(0, 1 - contrast), 1 + contrast].
+    saturation (float): How much to jitter saturation. saturation_factor
+    is chosen uniformly from [max(0, 1 - saturation), 1 + saturation].
+    hue(float): How much to jitter hue. hue_factor is chosen uniformly from
+    [-hue, hue]. Should be >=0 and <= 0.5.
+    """
+
+    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
+        self.brightness = brightness
+        self.contrast = contrast
+        self.saturation = saturation
+        self.hue = hue
+
+    def get_params(self, brightness, contrast, saturation, hue):
+        if brightness > 0:
+            brightness_factor = random.uniform(
+                max(0, 1 - brightness), 1 + brightness)
+        else:
+            brightness_factor = None
+
+        if contrast > 0:
+            contrast_factor = random.uniform(
+                max(0, 1 - contrast), 1 + contrast)
+        else:
+            contrast_factor = None
+
+        if saturation > 0:
+            saturation_factor = random.uniform(
+                max(0, 1 - saturation), 1 + saturation)
+        else:
+            saturation_factor = None
+
+        if hue > 0:
+            hue_factor = random.uniform(-hue, hue)
+        else:
+            hue_factor = None
+        return brightness_factor, contrast_factor, saturation_factor, hue_factor
+
+    def __call__(self, clip):
+        """
+        Args:
+        clip (list): list of PIL.Image
+        Returns:
+        list PIL.Image : list of transformed PIL.Image
+        """
+        if isinstance(clip[0], np.ndarray):
+            brightness, contrast, saturation, hue = self.get_params(
+                self.brightness, self.contrast, self.saturation, self.hue)
+
+            # Create img transform function sequence
+            img_transforms = []
+            if brightness is not None:
+                img_transforms.append(lambda img: torchvision.transforms.functional.adjust_brightness(img, brightness))
+            if saturation is not None:
+                img_transforms.append(lambda img: torchvision.transforms.functional.adjust_saturation(img, saturation))
+            if hue is not None:
+                img_transforms.append(lambda img: torchvision.transforms.functional.adjust_hue(img, hue))
+            if contrast is not None:
+                img_transforms.append(lambda img: torchvision.transforms.functional.adjust_contrast(img, contrast))
+            random.shuffle(img_transforms)
+            img_transforms = [img_as_ubyte, torchvision.transforms.ToPILImage()] + img_transforms + [np.array,
+                                                                                                     img_as_float]
+
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                jittered_clip = []
+                for img in clip:
+                    jittered_img = img
+                    for func in img_transforms:
+                        jittered_img = func(jittered_img)
+                    jittered_clip.append(jittered_img.astype('float32'))
+        elif isinstance(clip[0], PIL.Image.Image):
+            brightness, contrast, saturation, hue = self.get_params(
+                self.brightness, self.contrast, self.saturation, self.hue)
+
+            # Create img transform function sequence
+            img_transforms = []
+            if brightness is not None:
+                img_transforms.append(lambda img: torchvision.transforms.functional.adjust_brightness(img, brightness))
+            if saturation is not None:
+                img_transforms.append(lambda img: torchvision.transforms.functional.adjust_saturation(img, saturation))
+            if hue is not None:
+                img_transforms.append(lambda img: torchvision.transforms.functional.adjust_hue(img, hue))
+            if contrast is not None:
+                img_transforms.append(lambda img: torchvision.transforms.functional.adjust_contrast(img, contrast))
+            random.shuffle(img_transforms)
+
+            # Apply to all videos
+            jittered_clip = []
+            for img in clip:
+                for func in img_transforms:
+                    jittered_img = func(img)
+                jittered_clip.append(jittered_img)
+
+        else:
+            raise TypeError('Expected numpy.ndarray or PIL.Image' +
+                            'but got list of {0}'.format(type(clip[0])))
+        return jittered_clip
+
+
+class AllAugmentationTransform:
+    def __init__(self, resize_param=None, rotation_param=None, flip_param=None, crop_param=None, jitter_param=None):
+        self.transforms = []
+
+        if flip_param is not None:
+            self.transforms.append(RandomFlip(**flip_param))
+
+        if rotation_param is not None:
+            self.transforms.append(RandomRotation(**rotation_param))
+
+        if resize_param is not None:
+            self.transforms.append(RandomResize(**resize_param))
+
+        if crop_param is not None:
+            self.transforms.append(RandomCrop(**crop_param))
+
+        if jitter_param is not None:
+            self.transforms.append(ColorJitter(**jitter_param))
+
+    def __call__(self, clip):
+        for t in self.transforms:
+            clip = t(clip)
+        return clip

checkpoints/README.md

@@ -0,0 +1 @@
+# Checkpoints

config/mgif-256.yaml

@@ -0,0 +1,75 @@
+dataset_params:
+  root_dir: ../moving-gif
+  frame_shape: null
+  id_sampling: False
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    crop_param:
+      size: [256, 256]
+    resize_param:
+      ratio: [0.9, 1.1]
+    jitter_param:
+      hue: 0.5
+
+model_params:
+  common_params:
+    num_tps: 10
+    num_channels: 3
+    bg: False
+    multi_mask: True
+  generator_params:
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 3
+  dense_motion_params:
+    block_expansion: 64
+    max_features: 1024
+    num_blocks: 5
+    scale_factor: 0.25
+  avd_network_params:
+    id_bottle_size: 128
+    pose_bottle_size: 128
+
+
+train_params:
+  num_epochs: 100
+  num_repeats: 50
+  epoch_milestones: [70, 90]
+  lr_generator: 2.0e-4
+  batch_size: 28
+  scales: [1, 0.5, 0.25, 0.125]
+  dataloader_workers: 12
+  checkpoint_freq: 50
+  dropout_epoch: 35
+  dropout_maxp: 0.5
+  dropout_startp: 0.2
+  dropout_inc_epoch: 10
+  bg_start: 0
+  transform_params:
+    sigma_affine: 0.05
+    sigma_tps: 0.005
+    points_tps: 5
+  loss_weights:
+    perceptual: [10, 10, 10, 10, 10]
+    equivariance_value: 10
+    warp_loss: 10
+    bg: 10
+
+train_avd_params:
+  num_epochs: 100
+  num_repeats: 50
+  batch_size: 256
+  dataloader_workers: 24
+  checkpoint_freq: 10
+  epoch_milestones: [70, 90]
+  lr: 1.0e-3
+  lambda_shift: 1
+  lambda_affine: 1
+  random_scale: 0.25
+
+visualizer_params:
+  kp_size: 5
+  draw_border: True
+  colormap: 'gist_rainbow'

config/taichi-256.yaml

@@ -0,0 +1,134 @@
+# Dataset parameters
+# Each dataset should contain 2 folders train and test
+# Each video can be represented as:
+#   - an image of concatenated frames
+#   - '.mp4' or '.gif'
+#   - folder with all frames from a specific video
+# In case of Taichi. Same (youtube) video can be splitted in many parts (chunks). Each part has a following
+# format (id)#other#info.mp4. For example '12335#adsbf.mp4' has an id 12335. In case of TaiChi id stands for youtube
+# video id.
+dataset_params:
+  # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames.
+  root_dir: ../taichi
+  # Image shape, needed for staked .png format.
+  frame_shape: null
+  # In case of TaiChi single video can be splitted in many chunks, or the maybe several videos for single person.
+  # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False)
+  # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335
+  id_sampling: True
+  # Augmentation parameters see augmentation.py for all posible augmentations
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+# Defines model architecture
+model_params:
+  common_params:
+    # Number of TPS transformation
+    num_tps: 10
+    # Number of channels per image
+    num_channels: 3
+    # Whether to estimate affine background transformation
+    bg: True
+    # Whether to estimate the multi-resolution occlusion masks
+    multi_mask: True
+  generator_params:
+    # Number of features mutliplier
+    block_expansion: 64
+    # Maximum allowed number of features
+    max_features: 512
+    # Number of downsampling blocks and Upsampling blocks.
+    num_down_blocks: 3
+  dense_motion_params:
+    # Number of features mutliplier
+    block_expansion: 64
+    # Maximum allowed number of features
+    max_features: 1024
+    # Number of block in Unet. 
+    num_blocks: 5
+    # Optical flow is predicted on smaller images for better performance,
+    # scale_factor=0.25 means that 256x256 image will be resized to 64x64  
+    scale_factor: 0.25
+  avd_network_params:
+    # Bottleneck for identity branch
+    id_bottle_size: 128
+    # Bottleneck for pose branch
+    pose_bottle_size: 128
+
+# Parameters of training
+train_params:
+  # Number of training epochs 
+  num_epochs: 100
+  # For better i/o performance when number of videos is small number of epochs can be multiplied by this number.
+  # Thus effectivlly with num_repeats=100 each epoch is 100 times larger. 
+  num_repeats: 150
+  # Drop learning rate by 10 times after this epochs 
+  epoch_milestones: [70, 90]
+  # Initial learing rate for all modules
+  lr_generator: 2.0e-4
+  batch_size: 28
+  # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256,
+  # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32.
+  scales: [1, 0.5, 0.25, 0.125]
+  # Dataset preprocessing cpu workers
+  dataloader_workers: 12
+  # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs.
+  checkpoint_freq: 50
+  # Parameters of dropout
+  # The first dropout_epoch training uses dropout operation 
+  dropout_epoch: 35
+  # The probability P will linearly increase from dropout_startp to dropout_maxp in dropout_inc_epoch epochs
+  dropout_maxp: 0.7
+  dropout_startp: 0.0
+  dropout_inc_epoch: 10
+  # Estimate affine background transformation from the bg_start epoch.
+  bg_start: 0
+  # Parameters of random TPS transformation for equivariance loss
+  transform_params:
+    # Sigma for affine part
+    sigma_affine: 0.05
+    # Sigma for deformation part
+    sigma_tps: 0.005
+    # Number of point in the deformation grid
+    points_tps: 5
+  loss_weights:
+    # Weights for perceptual loss.
+    perceptual: [10, 10, 10, 10, 10]
+    # Weights for value equivariance.
+    equivariance_value: 10
+    # Weights for warp loss.
+    warp_loss: 10
+    # Weights for bg loss.
+    bg: 10
+
+# Parameters of training (animation-via-disentanglement)
+train_avd_params:
+  # Number of training epochs, visualization is produced after each epoch.
+  num_epochs: 100
+  # For better i/o performance when number of videos is small number of epochs can be multiplied by this number.
+  # Thus effectively with num_repeats=100 each epoch is 100 times larger.
+  num_repeats: 150
+  # Batch size.
+  batch_size: 256
+  # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs.
+  checkpoint_freq: 10
+  # Dataset preprocessing cpu workers
+  dataloader_workers: 24
+  # Drop learning rate 10 times after this epochs
+  epoch_milestones: [70, 90]
+  # Initial learning rate
+  lr: 1.0e-3
+  # Weights for equivariance loss.
+  lambda_shift: 1
+  random_scale: 0.25
+
+visualizer_params:
+  kp_size: 5
+  draw_border: True
+  colormap: 'gist_rainbow'

config/ted-384.yaml

@@ -0,0 +1,73 @@
+dataset_params:
+  root_dir: ../TED384-v2
+  frame_shape: null
+  id_sampling: True
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+model_params:
+  common_params:
+    num_tps: 10
+    num_channels: 3
+    bg: True
+    multi_mask: True
+  generator_params:
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 3
+  dense_motion_params:
+    block_expansion: 64
+    max_features: 1024
+    num_blocks: 5
+    scale_factor: 0.25
+  avd_network_params:
+    id_bottle_size: 128
+    pose_bottle_size: 128
+
+
+train_params:
+  num_epochs: 100
+  num_repeats: 150
+  epoch_milestones: [70, 90]
+  lr_generator: 2.0e-4
+  batch_size: 12
+  scales: [1, 0.5, 0.25, 0.125]
+  dataloader_workers: 6
+  checkpoint_freq: 50
+  dropout_epoch: 35
+  dropout_maxp: 0.5
+  dropout_startp: 0.0
+  dropout_inc_epoch: 10
+  bg_start: 0
+  transform_params:
+    sigma_affine: 0.05
+    sigma_tps: 0.005
+    points_tps: 5
+  loss_weights:
+    perceptual: [10, 10, 10, 10, 10]
+    equivariance_value: 10
+    warp_loss: 10
+    bg: 10
+
+train_avd_params:
+  num_epochs: 30
+  num_repeats: 500
+  batch_size: 256
+  dataloader_workers: 24
+  checkpoint_freq: 10
+  epoch_milestones: [20, 25]
+  lr: 1.0e-3
+  lambda_shift: 1
+  random_scale: 0.25
+
+visualizer_params:
+  kp_size: 5
+  draw_border: True
+  colormap: 'gist_rainbow'

config/vox-256.yaml

@@ -0,0 +1,74 @@
+dataset_params:
+  root_dir: ../vox
+  frame_shape: null
+  id_sampling: True
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+
+model_params:
+  common_params:
+    num_tps: 10
+    num_channels: 3
+    bg: True
+    multi_mask: True
+  generator_params:
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 3
+  dense_motion_params:
+    block_expansion: 64
+    max_features: 1024
+    num_blocks: 5
+    scale_factor: 0.25
+  avd_network_params:
+    id_bottle_size: 128
+    pose_bottle_size: 128
+
+
+train_params:
+  num_epochs: 100
+  num_repeats: 75
+  epoch_milestones: [70, 90]
+  lr_generator: 2.0e-4
+  batch_size: 28
+  scales: [1, 0.5, 0.25, 0.125]
+  dataloader_workers: 12
+  checkpoint_freq: 50
+  dropout_epoch: 35
+  dropout_maxp: 0.3
+  dropout_startp: 0.1
+  dropout_inc_epoch: 10
+  bg_start: 10
+  transform_params:
+    sigma_affine: 0.05
+    sigma_tps: 0.005
+    points_tps: 5
+  loss_weights:
+    perceptual: [10, 10, 10, 10, 10]
+    equivariance_value: 10
+    warp_loss: 10
+    bg: 10
+
+train_avd_params:
+  num_epochs: 200
+  num_repeats: 300
+  batch_size: 256
+  dataloader_workers: 24
+  checkpoint_freq: 50
+  epoch_milestones: [140, 180]
+  lr: 1.0e-3
+  lambda_shift: 1
+  random_scale: 0.25
+
+visualizer_params:
+  kp_size: 5
+  draw_border: True
+  colormap: 'gist_rainbow'

demo.py

@@ -0,0 +1,176 @@
+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
+
+if sys.version_info[0] < 3:
+    raise Exception("You must use Python 3 or higher. Recommended version is Python 3.9")
+
+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 load_checkpoints(config_path, checkpoint_path, device):
+    with open(config_path) as f:
+        config = yaml.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 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)):
+        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
+    return frame_num
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--config", required=True, help="path to config")
+    parser.add_argument("--checkpoint", default='checkpoints/vox.pth.tar', help="path to checkpoint to restore")
+
+    parser.add_argument("--source_image", default='./assets/source.png', help="path to source image")
+    parser.add_argument("--driving_video", default='./assets/driving.mp4', help="path to driving video")
+    parser.add_argument("--result_video", default='./result.mp4', help="path to output")
+    
+    parser.add_argument("--img_shape", default="256,256", type=lambda x: list(map(int, x.split(','))),
+                        help='Shape of image, that the model was trained on.')
+    
+    parser.add_argument("--mode", default='relative', choices=['standard', 'relative', 'avd'], help="Animate mode: ['standard', 'relative', 'avd'], when use the relative mode to animate a face, use '--find_best_frame' can get better quality result")
+    
+    parser.add_argument("--find_best_frame", dest="find_best_frame", action="store_true", 
+                        help="Generate from the frame that is the most alligned with source. (Only for faces, requires face_aligment lib)")
+
+    parser.add_argument("--cpu", dest="cpu", action="store_true", help="cpu mode.")
+
+    opt = parser.parse_args()
+
+    source_image = imageio.imread(opt.source_image)
+    reader = imageio.get_reader(opt.driving_video)
+    fps = reader.get_meta_data()['fps']
+    driving_video = []
+    try:
+        for im in reader:
+            driving_video.append(im)
+    except RuntimeError:
+        pass
+    reader.close()
+    
+    if opt.cpu:
+        device = torch.device('cpu')
+    else:
+        device = torch.device('cuda')
+    
+    source_image = resize(source_image, opt.img_shape)[..., :3]
+    driving_video = [resize(frame, opt.img_shape)[..., :3] for frame in driving_video]
+    inpainting, kp_detector, dense_motion_network, avd_network = load_checkpoints(config_path = opt.config, checkpoint_path = opt.checkpoint, device = device)
+ 
+    if opt.find_best_frame:
+        i = find_best_frame(source_image, driving_video, opt.cpu)
+        print ("Best frame: " + str(i))
+        driving_forward = driving_video[i:]
+        driving_backward = driving_video[:(i+1)][::-1]
+        predictions_forward = make_animation(source_image, driving_forward, inpainting, kp_detector, dense_motion_network, avd_network, device = device, mode = opt.mode)
+        predictions_backward = make_animation(source_image, driving_backward, inpainting, kp_detector, dense_motion_network, avd_network, device = device, mode = opt.mode)
+        predictions = predictions_backward[::-1] + predictions_forward[1:]
+    else:
+        predictions = make_animation(source_image, driving_video, inpainting, kp_detector, dense_motion_network, avd_network, device = device, mode = opt.mode)
+    
+    imageio.mimsave(opt.result_video, [img_as_ubyte(frame) for frame in predictions], fps=fps)
+

frames_dataset.py

@@ -0,0 +1,173 @@
+import os
+from skimage import io, img_as_float32
+from skimage.color import gray2rgb
+from sklearn.model_selection import train_test_split
+from imageio import mimread
+from skimage.transform import resize
+import numpy as np
+from torch.utils.data import Dataset
+from augmentation import AllAugmentationTransform
+import glob
+from functools import partial
+
+
+def read_video(name, frame_shape):
+    """
+    Read video which can be:
+      - an image of concatenated frames
+      - '.mp4' and'.gif'
+      - folder with videos
+    """
+
+    if os.path.isdir(name):
+        frames = sorted(os.listdir(name))
+        num_frames = len(frames)
+        video_array = np.array(
+            [img_as_float32(io.imread(os.path.join(name, frames[idx]))) for idx in range(num_frames)])
+    elif name.lower().endswith('.png') or name.lower().endswith('.jpg'):
+        image = io.imread(name)
+
+        if len(image.shape) == 2 or image.shape[2] == 1:
+            image = gray2rgb(image)
+
+        if image.shape[2] == 4:
+            image = image[..., :3]
+
+        image = img_as_float32(image)
+
+        video_array = np.moveaxis(image, 1, 0)
+
+        video_array = video_array.reshape((-1,) + frame_shape)
+        video_array = np.moveaxis(video_array, 1, 2)
+    elif name.lower().endswith('.gif') or name.lower().endswith('.mp4') or name.lower().endswith('.mov'):
+        video = mimread(name)
+        if len(video[0].shape) == 2:
+            video = [gray2rgb(frame) for frame in video]
+        if frame_shape is not None:
+            video = np.array([resize(frame, frame_shape) for frame in video])
+        video = np.array(video)
+        if video.shape[-1] == 4:
+            video = video[..., :3]
+        video_array = img_as_float32(video)
+    else:
+        raise Exception("Unknown file extensions  %s" % name)
+
+    return video_array
+
+
+class FramesDataset(Dataset):
+    """
+    Dataset of videos, each video can be represented as:
+      - an image of concatenated frames
+      - '.mp4' or '.gif'
+      - folder with all frames
+    """
+
+    def __init__(self, root_dir, frame_shape=(256, 256, 3), id_sampling=False, is_train=True,
+                 random_seed=0, pairs_list=None, augmentation_params=None):
+        self.root_dir = root_dir
+        self.videos = os.listdir(root_dir)
+        self.frame_shape = frame_shape
+        print(self.frame_shape)
+        self.pairs_list = pairs_list
+        self.id_sampling = id_sampling
+
+        if os.path.exists(os.path.join(root_dir, 'train')):
+            assert os.path.exists(os.path.join(root_dir, 'test'))
+            print("Use predefined train-test split.")
+            if id_sampling:
+                train_videos = {os.path.basename(video).split('#')[0] for video in
+                                os.listdir(os.path.join(root_dir, 'train'))}
+                train_videos = list(train_videos)
+            else:
+                train_videos = os.listdir(os.path.join(root_dir, 'train'))
+            test_videos = os.listdir(os.path.join(root_dir, 'test'))
+            self.root_dir = os.path.join(self.root_dir, 'train' if is_train else 'test')
+        else:
+            print("Use random train-test split.")
+            train_videos, test_videos = train_test_split(self.videos, random_state=random_seed, test_size=0.2)
+
+        if is_train:
+            self.videos = train_videos
+        else:
+            self.videos = test_videos
+
+        self.is_train = is_train
+
+        if self.is_train:
+            self.transform = AllAugmentationTransform(**augmentation_params)
+        else:
+            self.transform = None
+
+    def __len__(self):
+        return len(self.videos)
+
+    def __getitem__(self, idx):
+            
+        if self.is_train and self.id_sampling:   
+            name = self.videos[idx]
+            path = np.random.choice(glob.glob(os.path.join(self.root_dir, name + '*.mp4')))
+        else:
+            name = self.videos[idx]
+            path = os.path.join(self.root_dir, name)
+
+        video_name = os.path.basename(path)
+        if self.is_train and os.path.isdir(path):
+
+            frames = os.listdir(path)
+            num_frames = len(frames)
+            frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2))
+            
+            if self.frame_shape is not None:
+                resize_fn = partial(resize, output_shape=self.frame_shape)
+            else:
+                resize_fn = img_as_float32
+
+            if type(frames[0]) is bytes:
+                video_array = [resize_fn(io.imread(os.path.join(path, frames[idx].decode('utf-8')))) for idx in
+                            frame_idx]
+            else:
+                video_array = [resize_fn(io.imread(os.path.join(path, frames[idx]))) for idx in frame_idx]
+        else:
+                 
+            video_array = read_video(path, frame_shape=self.frame_shape)
+            
+            num_frames = len(video_array)
+            frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2)) if self.is_train else range(
+                num_frames)
+            video_array = video_array[frame_idx]
+            
+
+        if self.transform is not None:
+            video_array = self.transform(video_array)
+
+        out = {}
+        if self.is_train:
+            source = np.array(video_array[0], dtype='float32')
+            driving = np.array(video_array[1], dtype='float32')
+
+            out['driving'] = driving.transpose((2, 0, 1))
+            out['source'] = source.transpose((2, 0, 1))
+        else:
+            video = np.array(video_array, dtype='float32')
+            out['video'] = video.transpose((3, 0, 1, 2))
+
+        out['name'] = video_name
+        return out
+
+
+class DatasetRepeater(Dataset):
+    """
+    Pass several times over the same dataset for better i/o performance
+    """
+
+    def __init__(self, dataset, num_repeats=100):
+        self.dataset = dataset
+        self.num_repeats = num_repeats
+
+    def __len__(self):
+        return self.num_repeats * self.dataset.__len__()
+
+    def __getitem__(self, idx):
+        return self.dataset[idx % self.dataset.__len__()]
+

logger.py

@@ -0,0 +1,212 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import imageio
+
+import os
+from skimage.draw import circle
+
+import matplotlib.pyplot as plt
+import collections
+
+
+class Logger:
+    def __init__(self, log_dir, checkpoint_freq=50, visualizer_params=None, zfill_num=8, log_file_name='log.txt'):
+
+        self.loss_list = []
+        self.cpk_dir = log_dir
+        self.visualizations_dir = os.path.join(log_dir, 'train-vis')
+        if not os.path.exists(self.visualizations_dir):
+            os.makedirs(self.visualizations_dir)
+        self.log_file = open(os.path.join(log_dir, log_file_name), 'a')
+        self.zfill_num = zfill_num
+        self.visualizer = Visualizer(**visualizer_params)
+        self.checkpoint_freq = checkpoint_freq
+        self.epoch = 0
+        self.best_loss = float('inf')
+        self.names = None
+
+    def log_scores(self, loss_names):
+        loss_mean = np.array(self.loss_list).mean(axis=0)
+
+        loss_string = "; ".join(["%s - %.5f" % (name, value) for name, value in zip(loss_names, loss_mean)])
+        loss_string = str(self.epoch).zfill(self.zfill_num) + ") " + loss_string
+
+        print(loss_string, file=self.log_file)
+        self.loss_list = []
+        self.log_file.flush()
+
+    def visualize_rec(self, inp, out):
+        image = self.visualizer.visualize(inp['driving'], inp['source'], out)
+        imageio.imsave(os.path.join(self.visualizations_dir, "%s-rec.png" % str(self.epoch).zfill(self.zfill_num)), image)
+
+    def save_cpk(self, emergent=False):
+        cpk = {k: v.state_dict() for k, v in self.models.items()}
+        cpk['epoch'] = self.epoch
+        cpk_path = os.path.join(self.cpk_dir, '%s-checkpoint.pth.tar' % str(self.epoch).zfill(self.zfill_num)) 
+        if not (os.path.exists(cpk_path) and emergent):
+            torch.save(cpk, cpk_path)
+
+    @staticmethod
+    def load_cpk(checkpoint_path, inpainting_network=None, dense_motion_network =None, kp_detector=None, 
+                bg_predictor=None, avd_network=None, optimizer=None, optimizer_bg_predictor=None,
+                optimizer_avd=None):
+        checkpoint = torch.load(checkpoint_path)
+        if inpainting_network is not None:
+            inpainting_network.load_state_dict(checkpoint['inpainting_network'])
+        if kp_detector is not None:
+            kp_detector.load_state_dict(checkpoint['kp_detector'])
+        if bg_predictor is not None and 'bg_predictor' in checkpoint:
+            bg_predictor.load_state_dict(checkpoint['bg_predictor'])
+        if dense_motion_network is not None:
+            dense_motion_network.load_state_dict(checkpoint['dense_motion_network'])
+        if avd_network is not None:
+            if 'avd_network' in checkpoint:
+                avd_network.load_state_dict(checkpoint['avd_network'])
+        if optimizer_bg_predictor is not None and 'optimizer_bg_predictor' in checkpoint:
+            optimizer_bg_predictor.load_state_dict(checkpoint['optimizer_bg_predictor'])
+        if optimizer is not None and 'optimizer' in checkpoint:
+            optimizer.load_state_dict(checkpoint['optimizer'])
+        if optimizer_avd is not None:
+            if 'optimizer_avd' in checkpoint:
+                optimizer_avd.load_state_dict(checkpoint['optimizer_avd'])
+        epoch = -1
+        if 'epoch' in checkpoint:
+            epoch = checkpoint['epoch']
+        return epoch
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self):
+        if 'models' in self.__dict__:
+            self.save_cpk()
+        self.log_file.close()
+
+    def log_iter(self, losses):
+        losses = collections.OrderedDict(losses.items())
+        self.names = list(losses.keys())
+        self.loss_list.append(list(losses.values()))
+
+    def log_epoch(self, epoch, models, inp, out):
+        self.epoch = epoch
+        self.models = models
+        if (self.epoch + 1) % self.checkpoint_freq == 0:
+            self.save_cpk()
+        self.log_scores(self.names)
+        self.visualize_rec(inp, out)
+
+
+class Visualizer:
+    def __init__(self, kp_size=5, draw_border=False, colormap='gist_rainbow'):
+        self.kp_size = kp_size
+        self.draw_border = draw_border
+        self.colormap = plt.get_cmap(colormap)
+
+    def draw_image_with_kp(self, image, kp_array):
+        image = np.copy(image)
+        spatial_size = np.array(image.shape[:2][::-1])[np.newaxis]
+        kp_array = spatial_size * (kp_array + 1) / 2
+        num_kp = kp_array.shape[0]
+        for kp_ind, kp in enumerate(kp_array):
+            rr, cc = circle(kp[1], kp[0], self.kp_size, shape=image.shape[:2])
+            image[rr, cc] = np.array(self.colormap(kp_ind / num_kp))[:3]
+        return image
+
+    def create_image_column_with_kp(self, images, kp):
+        image_array = np.array([self.draw_image_with_kp(v, k) for v, k in zip(images, kp)])
+        return self.create_image_column(image_array)
+
+    def create_image_column(self, images):
+        if self.draw_border:
+            images = np.copy(images)
+            images[:, :, [0, -1]] = (1, 1, 1)
+            images[:, :, [0, -1]] = (1, 1, 1)
+        return np.concatenate(list(images), axis=0)
+
+    def create_image_grid(self, *args):
+        out = []
+        for arg in args:
+            if type(arg) == tuple:
+                out.append(self.create_image_column_with_kp(arg[0], arg[1]))
+            else:
+                out.append(self.create_image_column(arg))
+        return np.concatenate(out, axis=1)
+
+    def visualize(self, driving, source, out):
+        images = []
+
+        # Source image with keypoints
+        source = source.data.cpu()
+        kp_source = out['kp_source']['fg_kp'].data.cpu().numpy()
+        source = np.transpose(source, [0, 2, 3, 1])
+        images.append((source, kp_source))
+
+        # Equivariance visualization
+        if 'transformed_frame' in out:
+            transformed = out['transformed_frame'].data.cpu().numpy()
+            transformed = np.transpose(transformed, [0, 2, 3, 1])
+            transformed_kp = out['transformed_kp']['fg_kp'].data.cpu().numpy()
+            images.append((transformed, transformed_kp))
+
+        # Driving image with keypoints
+        kp_driving = out['kp_driving']['fg_kp'].data.cpu().numpy()
+        driving = driving.data.cpu().numpy()
+        driving = np.transpose(driving, [0, 2, 3, 1])
+        images.append((driving, kp_driving))
+
+        # Deformed image
+        if 'deformed' in out:
+            deformed = out['deformed'].data.cpu().numpy()
+            deformed = np.transpose(deformed, [0, 2, 3, 1])
+            images.append(deformed)
+
+        # Result with and without keypoints
+        prediction = out['prediction'].data.cpu().numpy()
+        prediction = np.transpose(prediction, [0, 2, 3, 1])
+        if 'kp_norm' in out:
+            kp_norm = out['kp_norm']['fg_kp'].data.cpu().numpy()
+            images.append((prediction, kp_norm))
+        images.append(prediction)
+
+
+        ## Occlusion map
+        if 'occlusion_map' in out:
+            for i in range(len(out['occlusion_map'])):
+                occlusion_map = out['occlusion_map'][i].data.cpu().repeat(1, 3, 1, 1)
+                occlusion_map = F.interpolate(occlusion_map, size=source.shape[1:3]).numpy()
+                occlusion_map = np.transpose(occlusion_map, [0, 2, 3, 1])
+                images.append(occlusion_map)
+
+        # Deformed images according to each individual transform
+        if 'deformed_source' in out:
+            full_mask = []
+            for i in range(out['deformed_source'].shape[1]):
+                image = out['deformed_source'][:, i].data.cpu()
+                # import ipdb;ipdb.set_trace()
+                image = F.interpolate(image, size=source.shape[1:3])
+                mask = out['contribution_maps'][:, i:(i+1)].data.cpu().repeat(1, 3, 1, 1)
+                mask = F.interpolate(mask, size=source.shape[1:3])
+                image = np.transpose(image.numpy(), (0, 2, 3, 1))
+                mask = np.transpose(mask.numpy(), (0, 2, 3, 1))
+
+                if i != 0:
+                    color = np.array(self.colormap((i - 1) / (out['deformed_source'].shape[1] - 1)))[:3]
+                else:
+                    color = np.array((0, 0, 0))
+
+                color = color.reshape((1, 1, 1, 3))
+
+                images.append(image)
+                if i != 0:
+                    images.append(mask * color)
+                else:
+                    images.append(mask)
+
+                full_mask.append(mask * color)
+
+            images.append(sum(full_mask))
+
+        image = self.create_image_grid(*images)
+        image = (255 * image).astype(np.uint8)
+        return image

modules/avd_network.py

@@ -0,0 +1,65 @@
+
+import torch
+from torch import nn
+
+
+class AVDNetwork(nn.Module):
+    """
+    Animation via Disentanglement network
+    """
+
+    def __init__(self, num_tps, id_bottle_size=64, pose_bottle_size=64):
+        super(AVDNetwork, self).__init__()
+        input_size = 5*2 * num_tps
+        self.num_tps = num_tps
+
+        self.id_encoder = nn.Sequential(
+            nn.Linear(input_size, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(inplace=True),
+            nn.Linear(256, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, 1024),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(inplace=True),
+            nn.Linear(1024, id_bottle_size)
+        )
+
+        self.pose_encoder = nn.Sequential(
+            nn.Linear(input_size, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(inplace=True),
+            nn.Linear(256, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, 1024),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(inplace=True),
+            nn.Linear(1024, pose_bottle_size)
+        )
+
+        self.decoder = nn.Sequential(
+            nn.Linear(pose_bottle_size + id_bottle_size, 1024),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(),
+            nn.Linear(1024, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(),
+            nn.Linear(256, input_size)
+        )
+
+    def forward(self, kp_source, kp_random):
+
+        bs = kp_source['fg_kp'].shape[0]
+        
+        pose_emb = self.pose_encoder(kp_random['fg_kp'].view(bs, -1))
+        id_emb = self.id_encoder(kp_source['fg_kp'].view(bs, -1))
+
+        rec = self.decoder(torch.cat([pose_emb, id_emb], dim=1))
+
+        rec =  {'fg_kp': rec.view(bs, self.num_tps*5, -1)}
+        return rec

modules/bg_motion_predictor.py

@@ -0,0 +1,24 @@
+from torch import nn
+import torch
+from torchvision import models
+
+class BGMotionPredictor(nn.Module):
+    """
+    Module for background estimation, return single transformation, parametrized as 3x3 matrix. The third row is [0 0 1]
+    """
+
+    def __init__(self):
+        super(BGMotionPredictor, self).__init__()
+        self.bg_encoder = models.resnet18(pretrained=False)
+        self.bg_encoder.conv1 = nn.Conv2d(6, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
+        num_features = self.bg_encoder.fc.in_features
+        self.bg_encoder.fc = nn.Linear(num_features, 6)
+        self.bg_encoder.fc.weight.data.zero_()
+        self.bg_encoder.fc.bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))
+
+    def forward(self, source_image, driving_image):
+        bs = source_image.shape[0]
+        out = torch.eye(3).unsqueeze(0).repeat(bs, 1, 1).type(source_image.type())
+        prediction = self.bg_encoder(torch.cat([source_image, driving_image], dim=1))
+        out[:, :2, :] = prediction.view(bs, 2, 3)
+        return out

modules/dense_motion.py

@@ -0,0 +1,164 @@
+from torch import nn
+import torch.nn.functional as F
+import torch
+from modules.util import Hourglass, AntiAliasInterpolation2d, make_coordinate_grid, kp2gaussian
+from modules.util import to_homogeneous, from_homogeneous, UpBlock2d, TPS
+import math
+
+class DenseMotionNetwork(nn.Module):
+    """
+    Module that estimating an optical flow and multi-resolution occlusion masks 
+                        from K TPS transformations and an affine transformation.
+    """
+
+    def __init__(self, block_expansion, num_blocks, max_features, num_tps, num_channels, 
+                 scale_factor=0.25, bg = False, multi_mask = True, kp_variance=0.01):
+        super(DenseMotionNetwork, self).__init__()
+
+        if scale_factor != 1:
+            self.down = AntiAliasInterpolation2d(num_channels, scale_factor)
+        self.scale_factor = scale_factor
+        self.multi_mask = multi_mask
+
+        self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_channels * (num_tps+1) + num_tps*5+1),
+                                   max_features=max_features, num_blocks=num_blocks)
+
+        hourglass_output_size = self.hourglass.out_channels
+        self.maps = nn.Conv2d(hourglass_output_size[-1], num_tps + 1, kernel_size=(7, 7), padding=(3, 3))
+
+        if multi_mask:
+            up = []
+            self.up_nums = int(math.log(1/scale_factor, 2))
+            self.occlusion_num = 4
+            
+            channel = [hourglass_output_size[-1]//(2**i) for i in range(self.up_nums)]
+            for i in range(self.up_nums):
+                up.append(UpBlock2d(channel[i], channel[i]//2, kernel_size=3, padding=1))
+            self.up = nn.ModuleList(up)
+
+            channel = [hourglass_output_size[-i-1] for i in range(self.occlusion_num-self.up_nums)[::-1]]
+            for i in range(self.up_nums):
+                channel.append(hourglass_output_size[-1]//(2**(i+1)))
+            occlusion = []
+            
+            for i in range(self.occlusion_num):
+                occlusion.append(nn.Conv2d(channel[i], 1, kernel_size=(7, 7), padding=(3, 3)))
+            self.occlusion = nn.ModuleList(occlusion)
+        else:
+            occlusion = [nn.Conv2d(hourglass_output_size[-1], 1, kernel_size=(7, 7), padding=(3, 3))]
+            self.occlusion = nn.ModuleList(occlusion)
+
+        self.num_tps = num_tps
+        self.bg = bg
+        self.kp_variance = kp_variance
+
+        
+    def create_heatmap_representations(self, source_image, kp_driving, kp_source):
+
+        spatial_size = source_image.shape[2:]
+        gaussian_driving = kp2gaussian(kp_driving['fg_kp'], spatial_size=spatial_size, kp_variance=self.kp_variance)
+        gaussian_source = kp2gaussian(kp_source['fg_kp'], spatial_size=spatial_size, kp_variance=self.kp_variance)
+        heatmap = gaussian_driving - gaussian_source
+
+        zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1]).type(heatmap.type()).to(heatmap.device)
+        heatmap = torch.cat([zeros, heatmap], dim=1)
+
+        return heatmap
+
+    def create_transformations(self, source_image, kp_driving, kp_source, bg_param):
+        # K TPS transformaions
+        bs, _, h, w = source_image.shape
+        kp_1 = kp_driving['fg_kp']
+        kp_2 = kp_source['fg_kp']
+        kp_1 = kp_1.view(bs, -1, 5, 2)
+        kp_2 = kp_2.view(bs, -1, 5, 2)
+        trans = TPS(mode = 'kp', bs = bs, kp_1 = kp_1, kp_2 = kp_2)
+        driving_to_source = trans.transform_frame(source_image)
+
+        identity_grid = make_coordinate_grid((h, w), type=kp_1.type()).to(kp_1.device)
+        identity_grid = identity_grid.view(1, 1, h, w, 2)
+        identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
+
+        # affine background transformation
+        if not (bg_param is None):            
+            identity_grid = to_homogeneous(identity_grid)
+            identity_grid = torch.matmul(bg_param.view(bs, 1, 1, 1, 3, 3), identity_grid.unsqueeze(-1)).squeeze(-1)
+            identity_grid = from_homogeneous(identity_grid)
+
+        transformations = torch.cat([identity_grid, driving_to_source], dim=1)
+        return transformations
+
+    def create_deformed_source_image(self, source_image, transformations):
+
+        bs, _, h, w = source_image.shape
+        source_repeat = source_image.unsqueeze(1).unsqueeze(1).repeat(1, self.num_tps + 1, 1, 1, 1, 1)
+        source_repeat = source_repeat.view(bs * (self.num_tps + 1), -1, h, w)
+        transformations = transformations.view((bs * (self.num_tps + 1), h, w, -1))
+        deformed = F.grid_sample(source_repeat, transformations, align_corners=True)
+        deformed = deformed.view((bs, self.num_tps+1, -1, h, w))
+        return deformed
+
+    def dropout_softmax(self, X, P):
+        '''
+        Dropout for TPS transformations. Eq(7) and Eq(8) in the paper.
+        '''
+        drop = (torch.rand(X.shape[0],X.shape[1]) < (1-P)).type(X.type()).to(X.device)
+        drop[..., 0] = 1
+        drop = drop.repeat(X.shape[2],X.shape[3],1,1).permute(2,3,0,1)
+
+        maxx = X.max(1).values.unsqueeze_(1)
+        X = X - maxx
+        X_exp = X.exp()
+        X[:,1:,...] /= (1-P)
+        mask_bool =(drop == 0)
+        X_exp = X_exp.masked_fill(mask_bool, 0)
+        partition = X_exp.sum(dim=1, keepdim=True) + 1e-6
+        return X_exp / partition  
+
+    def forward(self, source_image, kp_driving, kp_source, bg_param = None, dropout_flag=False, dropout_p = 0):
+        if self.scale_factor != 1:
+            source_image = self.down(source_image)
+
+        bs, _, h, w = source_image.shape
+
+        out_dict = dict()
+        heatmap_representation = self.create_heatmap_representations(source_image, kp_driving, kp_source)
+        transformations = self.create_transformations(source_image, kp_driving, kp_source, bg_param)
+        deformed_source = self.create_deformed_source_image(source_image, transformations)
+        out_dict['deformed_source'] = deformed_source
+        # out_dict['transformations'] = transformations
+        deformed_source = deformed_source.view(bs,-1,h,w)
+        input = torch.cat([heatmap_representation, deformed_source], dim=1)
+        input = input.view(bs, -1, h, w)
+
+        prediction = self.hourglass(input, mode = 1)
+
+        contribution_maps = self.maps(prediction[-1]) 
+        if(dropout_flag):
+            contribution_maps = self.dropout_softmax(contribution_maps, dropout_p)
+        else:
+            contribution_maps = F.softmax(contribution_maps, dim=1)
+        out_dict['contribution_maps'] = contribution_maps
+
+        # Combine the K+1 transformations
+        # Eq(6) in the paper
+        contribution_maps = contribution_maps.unsqueeze(2)
+        transformations = transformations.permute(0, 1, 4, 2, 3)
+        deformation = (transformations * contribution_maps).sum(dim=1)
+        deformation = deformation.permute(0, 2, 3, 1)
+
+        out_dict['deformation'] = deformation # Optical Flow
+
+        occlusion_map = []
+        if self.multi_mask:
+            for i in range(self.occlusion_num-self.up_nums):
+                occlusion_map.append(torch.sigmoid(self.occlusion[i](prediction[self.up_nums-self.occlusion_num+i])))
+            prediction = prediction[-1]
+            for i in range(self.up_nums):
+                prediction = self.up[i](prediction)
+                occlusion_map.append(torch.sigmoid(self.occlusion[i+self.occlusion_num-self.up_nums](prediction)))
+        else:
+            occlusion_map.append(torch.sigmoid(self.occlusion[0](prediction[-1])))
+                
+        out_dict['occlusion_map'] = occlusion_map # Multi-resolution Occlusion Masks
+        return out_dict

modules/inpainting_network.py

@@ -0,0 +1,130 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from modules.util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d
+from modules.dense_motion import DenseMotionNetwork
+
+
+class InpaintingNetwork(nn.Module):
+    """
+    Inpaint the missing regions and reconstruct the Driving image.
+    """
+    def __init__(self, num_channels, block_expansion, max_features, num_down_blocks, multi_mask = True, **kwargs):
+        super(InpaintingNetwork, self).__init__()
+
+        self.num_down_blocks = num_down_blocks
+        self.multi_mask = multi_mask
+        self.first = SameBlock2d(num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3))
+
+        down_blocks = []
+        for i in range(num_down_blocks):
+            in_features = min(max_features, block_expansion * (2 ** i))
+            out_features = min(max_features, block_expansion * (2 ** (i + 1)))
+            down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))
+        self.down_blocks = nn.ModuleList(down_blocks)
+
+        up_blocks = []
+        in_features = [max_features, max_features, max_features//2]
+        out_features = [max_features//2, max_features//4, max_features//8]
+        for i in range(num_down_blocks):
+            up_blocks.append(UpBlock2d(in_features[i], out_features[i], kernel_size=(3, 3), padding=(1, 1)))
+        self.up_blocks = nn.ModuleList(up_blocks)
+
+        resblock = []
+        for i in range(num_down_blocks):
+            resblock.append(ResBlock2d(in_features[i], kernel_size=(3, 3), padding=(1, 1)))
+            resblock.append(ResBlock2d(in_features[i], kernel_size=(3, 3), padding=(1, 1)))
+        self.resblock = nn.ModuleList(resblock)
+
+        self.final = nn.Conv2d(block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3))
+        self.num_channels = num_channels
+
+    def deform_input(self, inp, deformation):
+        _, h_old, w_old, _ = deformation.shape
+        _, _, h, w = inp.shape
+        if h_old != h or w_old != w:
+            deformation = deformation.permute(0, 3, 1, 2)
+            deformation = F.interpolate(deformation, size=(h, w), mode='bilinear', align_corners=True)
+            deformation = deformation.permute(0, 2, 3, 1)
+        return F.grid_sample(inp, deformation,align_corners=True)
+
+    def occlude_input(self, inp, occlusion_map):
+        if not self.multi_mask:
+            if inp.shape[2] != occlusion_map.shape[2] or inp.shape[3] != occlusion_map.shape[3]:
+                occlusion_map = F.interpolate(occlusion_map, size=inp.shape[2:], mode='bilinear',align_corners=True)
+        out = inp * occlusion_map
+        return out
+
+    def forward(self, source_image, dense_motion):
+        out = self.first(source_image) 
+        encoder_map = [out]
+        for i in range(len(self.down_blocks)):
+            out = self.down_blocks[i](out)
+            encoder_map.append(out)
+
+        output_dict = {}
+        output_dict['contribution_maps'] = dense_motion['contribution_maps']
+        output_dict['deformed_source'] = dense_motion['deformed_source']
+
+        occlusion_map = dense_motion['occlusion_map']
+        output_dict['occlusion_map'] = occlusion_map
+
+        deformation = dense_motion['deformation']
+        out_ij = self.deform_input(out.detach(), deformation)
+        out = self.deform_input(out, deformation)
+
+        out_ij = self.occlude_input(out_ij, occlusion_map[0].detach())
+        out = self.occlude_input(out, occlusion_map[0])
+
+        warped_encoder_maps = []
+        warped_encoder_maps.append(out_ij)
+
+        for i in range(self.num_down_blocks):
+            
+            out = self.resblock[2*i](out)
+            out = self.resblock[2*i+1](out)
+            out = self.up_blocks[i](out)
+            
+            encode_i = encoder_map[-(i+2)]
+            encode_ij = self.deform_input(encode_i.detach(), deformation)
+            encode_i = self.deform_input(encode_i, deformation)
+            
+            occlusion_ind = 0
+            if self.multi_mask:
+                occlusion_ind = i+1
+            encode_ij = self.occlude_input(encode_ij, occlusion_map[occlusion_ind].detach())
+            encode_i = self.occlude_input(encode_i, occlusion_map[occlusion_ind])
+            warped_encoder_maps.append(encode_ij)
+
+            if(i==self.num_down_blocks-1):
+                break
+
+            out = torch.cat([out, encode_i], 1)
+
+        deformed_source = self.deform_input(source_image, deformation)
+        output_dict["deformed"] = deformed_source
+        output_dict["warped_encoder_maps"] = warped_encoder_maps
+
+        occlusion_last = occlusion_map[-1]
+        if not self.multi_mask:
+            occlusion_last = F.interpolate(occlusion_last, size=out.shape[2:], mode='bilinear',align_corners=True)
+
+        out = out * (1 - occlusion_last) + encode_i
+        out = self.final(out)
+        out = torch.sigmoid(out)
+        out = out * (1 - occlusion_last) + deformed_source * occlusion_last
+        output_dict["prediction"] = out
+
+        return output_dict
+
+    def get_encode(self, driver_image, occlusion_map):
+        out = self.first(driver_image)
+        encoder_map = []
+        encoder_map.append(self.occlude_input(out.detach(), occlusion_map[-1].detach()))
+        for i in range(len(self.down_blocks)):
+            out = self.down_blocks[i](out.detach())
+            out_mask = self.occlude_input(out.detach(), occlusion_map[2-i].detach())
+            encoder_map.append(out_mask.detach())
+
+        return encoder_map
+

modules/keypoint_detector.py

@@ -0,0 +1,27 @@
+from torch import nn
+import torch
+from torchvision import models
+
+class KPDetector(nn.Module):
+    """
+    Predict K*5 keypoints.
+    """
+
+    def __init__(self, num_tps, **kwargs):
+        super(KPDetector, self).__init__()
+        self.num_tps = num_tps
+
+        self.fg_encoder = models.resnet18(pretrained=False)
+        num_features = self.fg_encoder.fc.in_features
+        self.fg_encoder.fc = nn.Linear(num_features, num_tps*5*2)
+
+        
+    def forward(self, image):
+
+        fg_kp = self.fg_encoder(image)
+        bs, _, = fg_kp.shape
+        fg_kp = torch.sigmoid(fg_kp)
+        fg_kp = fg_kp * 2 - 1
+        out = {'fg_kp': fg_kp.view(bs, self.num_tps*5, -1)}
+
+        return out

modules/model.py

@@ -0,0 +1,182 @@
+from torch import nn
+import torch
+import torch.nn.functional as F
+from modules.util import AntiAliasInterpolation2d, TPS
+from torchvision import models
+import numpy as np
+
+
+class Vgg19(torch.nn.Module):
+    """
+    Vgg19 network for perceptual loss. See Sec 3.3.
+    """
+    def __init__(self, requires_grad=False):
+        super(Vgg19, self).__init__()
+        vgg_pretrained_features = models.vgg19(pretrained=True).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        for x in range(2):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+
+        self.mean = torch.nn.Parameter(data=torch.Tensor(np.array([0.485, 0.456, 0.406]).reshape((1, 3, 1, 1))),
+                                       requires_grad=False)
+        self.std = torch.nn.Parameter(data=torch.Tensor(np.array([0.229, 0.224, 0.225]).reshape((1, 3, 1, 1))),
+                                      requires_grad=False)
+
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        X = (X - self.mean) / self.std
+        h_relu1 = self.slice1(X)
+        h_relu2 = self.slice2(h_relu1)
+        h_relu3 = self.slice3(h_relu2)
+        h_relu4 = self.slice4(h_relu3)
+        h_relu5 = self.slice5(h_relu4)
+        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+        return out
+
+
+class ImagePyramide(torch.nn.Module):
+    """
+    Create image pyramide for computing pyramide perceptual loss. See Sec 3.3
+    """
+    def __init__(self, scales, num_channels):
+        super(ImagePyramide, self).__init__()
+        downs = {}
+        for scale in scales:
+            downs[str(scale).replace('.', '-')] = AntiAliasInterpolation2d(num_channels, scale)
+        self.downs = nn.ModuleDict(downs)
+
+    def forward(self, x):
+        out_dict = {}
+        for scale, down_module in self.downs.items():
+            out_dict['prediction_' + str(scale).replace('-', '.')] = down_module(x)
+        return out_dict
+
+
+def detach_kp(kp):
+    return {key: value.detach() for key, value in kp.items()}
+
+
+class GeneratorFullModel(torch.nn.Module):
+    """
+    Merge all generator related updates into single model for better multi-gpu usage
+    """
+
+    def __init__(self, kp_extractor, bg_predictor, dense_motion_network, inpainting_network, train_params, *kwargs):
+        super(GeneratorFullModel, self).__init__()
+        self.kp_extractor = kp_extractor
+        self.inpainting_network = inpainting_network
+        self.dense_motion_network = dense_motion_network
+
+        self.bg_predictor = None
+        if bg_predictor:
+            self.bg_predictor = bg_predictor
+            self.bg_start = train_params['bg_start']
+
+        self.train_params = train_params
+        self.scales = train_params['scales']
+
+        self.pyramid = ImagePyramide(self.scales, inpainting_network.num_channels)
+        if torch.cuda.is_available():
+            self.pyramid = self.pyramid.cuda()
+
+        self.loss_weights = train_params['loss_weights']
+        self.dropout_epoch = train_params['dropout_epoch']
+        self.dropout_maxp = train_params['dropout_maxp']
+        self.dropout_inc_epoch = train_params['dropout_inc_epoch']
+        self.dropout_startp =train_params['dropout_startp']
+        
+        if sum(self.loss_weights['perceptual']) != 0:
+            self.vgg = Vgg19()
+            if torch.cuda.is_available():
+                self.vgg = self.vgg.cuda()
+
+
+    def forward(self, x, epoch):
+        kp_source = self.kp_extractor(x['source'])
+        kp_driving = self.kp_extractor(x['driving'])
+        bg_param = None
+        if self.bg_predictor:
+            if(epoch>=self.bg_start):
+                bg_param = self.bg_predictor(x['source'], x['driving'])
+          
+        if(epoch>=self.dropout_epoch):
+            dropout_flag = False
+            dropout_p = 0
+        else:
+            # dropout_p will linearly increase from dropout_startp to dropout_maxp 
+            dropout_flag = True
+            dropout_p = min(epoch/self.dropout_inc_epoch * self.dropout_maxp + self.dropout_startp, self.dropout_maxp)
+        
+        dense_motion = self.dense_motion_network(source_image=x['source'], kp_driving=kp_driving,
+                                                    kp_source=kp_source, bg_param = bg_param, 
+                                                    dropout_flag = dropout_flag, dropout_p = dropout_p)
+        generated = self.inpainting_network(x['source'], dense_motion)
+        generated.update({'kp_source': kp_source, 'kp_driving': kp_driving})
+
+        loss_values = {}
+
+        pyramide_real = self.pyramid(x['driving'])
+        pyramide_generated = self.pyramid(generated['prediction'])
+
+        # reconstruction loss
+        if sum(self.loss_weights['perceptual']) != 0:
+            value_total = 0
+            for scale in self.scales:
+                x_vgg = self.vgg(pyramide_generated['prediction_' + str(scale)])
+                y_vgg = self.vgg(pyramide_real['prediction_' + str(scale)])
+
+                for i, weight in enumerate(self.loss_weights['perceptual']):
+                    value = torch.abs(x_vgg[i] - y_vgg[i].detach()).mean()
+                    value_total += self.loss_weights['perceptual'][i] * value
+            loss_values['perceptual'] = value_total
+
+        # equivariance loss
+        if self.loss_weights['equivariance_value'] != 0:
+            transform_random = TPS(mode = 'random', bs = x['driving'].shape[0], **self.train_params['transform_params'])
+            transform_grid = transform_random.transform_frame(x['driving'])
+            transformed_frame = F.grid_sample(x['driving'], transform_grid, padding_mode="reflection",align_corners=True)
+            transformed_kp = self.kp_extractor(transformed_frame)
+
+            generated['transformed_frame'] = transformed_frame
+            generated['transformed_kp'] = transformed_kp
+        
+            warped = transform_random.warp_coordinates(transformed_kp['fg_kp'])
+            kp_d = kp_driving['fg_kp']
+            value = torch.abs(kp_d - warped).mean()
+            loss_values['equivariance_value'] = self.loss_weights['equivariance_value'] * value
+
+        # warp loss
+        if self.loss_weights['warp_loss'] != 0:
+            occlusion_map = generated['occlusion_map']
+            encode_map = self.inpainting_network.get_encode(x['driving'], occlusion_map)
+            decode_map = generated['warped_encoder_maps']
+            value = 0
+            for i in range(len(encode_map)):
+                value += torch.abs(encode_map[i]-decode_map[-i-1]).mean()
+
+            loss_values['warp_loss'] = self.loss_weights['warp_loss'] * value
+        
+        # bg loss
+        if self.bg_predictor and epoch >= self.bg_start and self.loss_weights['bg'] != 0:
+            bg_param_reverse = self.bg_predictor(x['driving'], x['source'])
+            value = torch.matmul(bg_param, bg_param_reverse)
+            eye = torch.eye(3).view(1, 1, 3, 3).type(value.type())
+            value = torch.abs(eye - value).mean()
+            loss_values['bg'] = self.loss_weights['bg'] * value
+
+        return loss_values, generated

modules/util.py

@@ -0,0 +1,349 @@
+from torch import nn
+import torch.nn.functional as F
+import torch
+
+
+class TPS:
+    '''
+    TPS transformation, mode 'kp' for Eq(2) in the paper, mode 'random' for equivariance loss.
+    '''
+    def __init__(self, mode, bs, **kwargs):
+        self.bs = bs
+        self.mode = mode
+        if mode == 'random':
+            noise = torch.normal(mean=0, std=kwargs['sigma_affine'] * torch.ones([bs, 2, 3]))
+            self.theta = noise + torch.eye(2, 3).view(1, 2, 3)
+            self.control_points = make_coordinate_grid((kwargs['points_tps'], kwargs['points_tps']), type=noise.type())
+            self.control_points = self.control_points.unsqueeze(0)
+            self.control_params = torch.normal(mean=0, 
+                        std=kwargs['sigma_tps'] * torch.ones([bs, 1, kwargs['points_tps'] ** 2]))
+        elif mode == 'kp':
+            kp_1 = kwargs["kp_1"]
+            kp_2 = kwargs["kp_2"]
+            device = kp_1.device
+            kp_type = kp_1.type()
+            self.gs = kp_1.shape[1]
+            n = kp_1.shape[2]
+            K = torch.norm(kp_1[:,:,:, None]-kp_1[:,:, None, :], dim=4, p=2)
+            K = K**2
+            K = K * torch.log(K+1e-9)
+            
+            one1 = torch.ones(self.bs, kp_1.shape[1], kp_1.shape[2], 1).to(device).type(kp_type)
+            kp_1p = torch.cat([kp_1,one1], 3)
+            
+            zero = torch.zeros(self.bs, kp_1.shape[1], 3, 3).to(device).type(kp_type)
+            P = torch.cat([kp_1p, zero],2)
+            L = torch.cat([K,kp_1p.permute(0,1,3,2)],2)
+            L = torch.cat([L,P],3)
+        
+            zero = torch.zeros(self.bs, kp_1.shape[1], 3, 2).to(device).type(kp_type)
+            Y = torch.cat([kp_2, zero], 2)
+            one = torch.eye(L.shape[2]).expand(L.shape).to(device).type(kp_type)*0.01
+            L = L + one
+
+            param = torch.matmul(torch.inverse(L),Y)
+            self.theta = param[:,:,n:,:].permute(0,1,3,2)
+
+            self.control_points = kp_1
+            self.control_params = param[:,:,:n,:]
+        else:
+            raise Exception("Error TPS mode")
+
+    def transform_frame(self, frame):
+        grid = make_coordinate_grid(frame.shape[2:], type=frame.type()).unsqueeze(0).to(frame.device)
+        grid = grid.view(1, frame.shape[2] * frame.shape[3], 2)
+        shape = [self.bs, frame.shape[2], frame.shape[3], 2]
+        if self.mode == 'kp':
+            shape.insert(1, self.gs)
+        grid = self.warp_coordinates(grid).view(*shape)
+        return grid
+
+    def warp_coordinates(self, coordinates):
+        theta = self.theta.type(coordinates.type()).to(coordinates.device)
+        control_points = self.control_points.type(coordinates.type()).to(coordinates.device)
+        control_params = self.control_params.type(coordinates.type()).to(coordinates.device)
+
+        if self.mode == 'kp':
+            transformed = torch.matmul(theta[:, :, :, :2], coordinates.permute(0, 2, 1)) + theta[:, :, :, 2:]
+
+            distances = coordinates.view(coordinates.shape[0], 1, 1, -1, 2) - control_points.view(self.bs, control_points.shape[1], -1, 1, 2)
+
+            distances = distances ** 2
+            result = distances.sum(-1)
+            result = result * torch.log(result + 1e-9)
+            result = torch.matmul(result.permute(0, 1, 3, 2), control_params)
+            transformed = transformed.permute(0, 1, 3, 2) + result
+
+        elif self.mode == 'random':
+            theta = theta.unsqueeze(1)
+            transformed = torch.matmul(theta[:, :, :, :2], coordinates.unsqueeze(-1)) + theta[:, :, :, 2:]
+            transformed = transformed.squeeze(-1)
+            ances = coordinates.view(coordinates.shape[0], -1, 1, 2) - control_points.view(1, 1, -1, 2)
+            distances = ances ** 2
+
+            result = distances.sum(-1)
+            result = result * torch.log(result + 1e-9)
+            result = result * control_params
+            result = result.sum(dim=2).view(self.bs, coordinates.shape[1], 1)
+            transformed = transformed + result
+        else:
+            raise Exception("Error TPS mode")
+
+        return transformed
+        
+
+def kp2gaussian(kp, spatial_size, kp_variance):
+    """
+    Transform a keypoint into gaussian like representation
+    """
+
+    coordinate_grid = make_coordinate_grid(spatial_size, kp.type()).to(kp.device)
+    number_of_leading_dimensions = len(kp.shape) - 1
+    shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape
+    coordinate_grid = coordinate_grid.view(*shape)
+    repeats = kp.shape[:number_of_leading_dimensions] + (1, 1, 1)
+    coordinate_grid = coordinate_grid.repeat(*repeats)
+
+    # Preprocess kp shape
+    shape = kp.shape[:number_of_leading_dimensions] + (1, 1, 2)
+    kp = kp.view(*shape)
+
+    mean_sub = (coordinate_grid - kp)
+
+    out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance)
+
+    return out
+
+
+def make_coordinate_grid(spatial_size, type):
+    """
+    Create a meshgrid [-1,1] x [-1,1] of given spatial_size.
+    """
+    h, w = spatial_size
+    x = torch.arange(w).type(type)
+    y = torch.arange(h).type(type)
+
+    x = (2 * (x / (w - 1)) - 1)
+    y = (2 * (y / (h - 1)) - 1)
+
+    yy = y.view(-1, 1).repeat(1, w)
+    xx = x.view(1, -1).repeat(h, 1)
+
+    meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2)
+
+    return meshed
+
+
+class ResBlock2d(nn.Module):
+    """
+    Res block, preserve spatial resolution.
+    """
+
+    def __init__(self, in_features, kernel_size, padding):
+        super(ResBlock2d, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size,
+                               padding=padding)
+        self.conv2 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size,
+                               padding=padding)
+        self.norm1 = nn.InstanceNorm2d(in_features, affine=True)
+        self.norm2 = nn.InstanceNorm2d(in_features, affine=True)
+
+    def forward(self, x):
+        out = self.norm1(x)
+        out = F.relu(out)
+        out = self.conv1(out)
+        out = self.norm2(out)
+        out = F.relu(out)
+        out = self.conv2(out)
+        out += x
+        return out
+
+
+class UpBlock2d(nn.Module):
+    """
+    Upsampling block for use in decoder.
+    """
+
+    def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
+        super(UpBlock2d, self).__init__()
+
+        self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
+                              padding=padding, groups=groups)
+        self.norm = nn.InstanceNorm2d(out_features, affine=True)
+
+    def forward(self, x):
+        out = F.interpolate(x, scale_factor=2)
+        out = self.conv(out)
+        out = self.norm(out)
+        out = F.relu(out)
+        return out
+
+
+class DownBlock2d(nn.Module):
+    """
+    Downsampling block for use in encoder.
+    """
+
+    def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
+        super(DownBlock2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
+                              padding=padding, groups=groups)
+        self.norm = nn.InstanceNorm2d(out_features, affine=True)
+        self.pool = nn.AvgPool2d(kernel_size=(2, 2))
+
+    def forward(self, x):
+        out = self.conv(x)
+        out = self.norm(out)
+        out = F.relu(out)
+        out = self.pool(out)
+        return out
+
+
+class SameBlock2d(nn.Module):
+    """
+    Simple block, preserve spatial resolution.
+    """
+
+    def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1):
+        super(SameBlock2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features,
+                              kernel_size=kernel_size, padding=padding, groups=groups)
+        self.norm = nn.InstanceNorm2d(out_features, affine=True)
+
+    def forward(self, x):
+        out = self.conv(x)
+        out = self.norm(out)
+        out = F.relu(out)
+        return out
+
+
+class Encoder(nn.Module):
+    """
+    Hourglass Encoder
+    """
+
+    def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
+        super(Encoder, self).__init__()
+
+        down_blocks = []
+        for i in range(num_blocks):
+            down_blocks.append(DownBlock2d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)),
+                                           min(max_features, block_expansion * (2 ** (i + 1))),
+                                           kernel_size=3, padding=1))
+        self.down_blocks = nn.ModuleList(down_blocks)
+
+    def forward(self, x):
+        outs = [x]
+        #print('encoder:' ,outs[-1].shape)
+        for down_block in self.down_blocks:
+            outs.append(down_block(outs[-1]))
+            #print('encoder:' ,outs[-1].shape)
+        return outs
+
+
+class Decoder(nn.Module):
+    """
+    Hourglass Decoder
+    """
+
+    def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
+        super(Decoder, self).__init__()
+
+        up_blocks = []
+        self.out_channels = []
+        for i in range(num_blocks)[::-1]:
+            in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1)))
+            self.out_channels.append(in_filters)
+            out_filters = min(max_features, block_expansion * (2 ** i))
+            up_blocks.append(UpBlock2d(in_filters, out_filters, kernel_size=3, padding=1))
+
+        self.up_blocks = nn.ModuleList(up_blocks)
+        self.out_channels.append(block_expansion + in_features)
+        # self.out_filters = block_expansion + in_features
+
+    def forward(self, x, mode = 0):
+        out = x.pop()
+        outs = []
+        for up_block in self.up_blocks:
+            out = up_block(out)
+            skip = x.pop()
+            out = torch.cat([out, skip], dim=1)
+            outs.append(out)
+        if(mode == 0):
+            return out
+        else:
+            return outs
+
+
+class Hourglass(nn.Module):
+    """
+    Hourglass architecture.
+    """
+
+    def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
+        super(Hourglass, self).__init__()
+        self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features)
+        self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features)
+        self.out_channels = self.decoder.out_channels
+        # self.out_filters = self.decoder.out_filters
+
+    def forward(self, x, mode = 0):
+        return self.decoder(self.encoder(x), mode)
+
+
+class AntiAliasInterpolation2d(nn.Module):
+    """
+    Band-limited downsampling, for better preservation of the input signal.
+    """
+    def __init__(self, channels, scale):
+        super(AntiAliasInterpolation2d, self).__init__()
+        sigma = (1 / scale - 1) / 2
+        kernel_size = 2 * round(sigma * 4) + 1
+        self.ka = kernel_size // 2
+        self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka
+
+        kernel_size = [kernel_size, kernel_size]
+        sigma = [sigma, sigma]
+        # The gaussian kernel is the product of the
+        # gaussian function of each dimension.
+        kernel = 1
+        meshgrids = torch.meshgrid(
+            [
+                torch.arange(size, dtype=torch.float32)
+                for size in kernel_size
+                ]
+        )
+        for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
+            mean = (size - 1) / 2
+            kernel *= torch.exp(-(mgrid - mean) ** 2 / (2 * std ** 2))
+
+        # Make sure sum of values in gaussian kernel equals 1.
+        kernel = kernel / torch.sum(kernel)
+        # Reshape to depthwise convolutional weight
+        kernel = kernel.view(1, 1, *kernel.size())
+        kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))
+
+        self.register_buffer('weight', kernel)
+        self.groups = channels
+        self.scale = scale
+
+    def forward(self, input):
+        if self.scale == 1.0:
+            return input
+
+        out = F.pad(input, (self.ka, self.kb, self.ka, self.kb))
+        out = F.conv2d(out, weight=self.weight, groups=self.groups)
+        out = F.interpolate(out, scale_factor=(self.scale, self.scale))
+
+        return out
+
+
+def to_homogeneous(coordinates):
+    ones_shape = list(coordinates.shape)
+    ones_shape[-1] = 1
+    ones = torch.ones(ones_shape).type(coordinates.type())
+
+    return torch.cat([coordinates, ones], dim=-1)
+
+def from_homogeneous(coordinates):
+    return coordinates[..., :2] / coordinates[..., 2:3]

reconstruction.py

@@ -0,0 +1,69 @@
+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
+
+
+def reconstruction(config, inpainting_network, kp_detector, bg_predictor, dense_motion_network, 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, inpainting_network=inpainting_network, kp_detector=kp_detector,
+                         bg_predictor=bg_predictor, dense_motion_network=dense_motion_network)
+    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 = []
+
+    inpainting_network.eval()
+    kp_detector.eval()
+    dense_motion_network.eval()
+    if bg_predictor:
+        bg_predictor.eval()
+
+    for it, x in tqdm(enumerate(dataloader)):
+        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)
+                bg_params = None
+                if bg_predictor:
+                    bg_params = bg_predictor(source, driving)
+                
+                dense_motion = dense_motion_network(source_image=source, kp_driving=kp_driving,
+                                                    kp_source=kp_source, bg_param = bg_params, 
+                                                    dropout_flag = False)
+                out = inpainting_network(source, dense_motion)
+                out['kp_source'] = kp_source
+                out['kp_driving'] = kp_driving
+
+                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 = torch.abs(out['prediction'] - driving).mean().cpu().numpy()
+                
+                loss_list.append(loss)
+            # print(np.mean(loss_list))
+            predictions = np.concatenate(predictions, axis=1)
+            imageio.imsave(os.path.join(png_dir, x['name'][0] + '.png'), (255 * predictions).astype(np.uint8))
+
+    print("Reconstruction loss: %s" % np.mean(loss_list))

requirements.txt

@@ -0,0 +1,26 @@
+cffi==1.14.6
+cycler==0.10.0
+decorator==5.1.0
+face-alignment==1.3.5
+imageio==2.9.0
+imageio-ffmpeg==0.4.5
+kiwisolver==1.3.2
+matplotlib==3.4.3
+networkx==2.6.3
+numpy==1.20.3
+pandas==1.3.3
+Pillow==8.3.2
+pycparser==2.20
+pyparsing==2.4.7
+python-dateutil==2.8.2
+pytz==2021.1
+PyWavelets==1.1.1
+PyYAML==5.4.1
+scikit-image==0.18.3
+scikit-learn==1.0
+scipy==1.7.1
+six==1.16.0
+torch==1.11.0
+torchvision==0.12.0
+tqdm==4.62.3
+gradio

run.py

@@ -0,0 +1,89 @@
+import matplotlib
+matplotlib.use('Agg')
+
+import os, sys
+import yaml
+from argparse import ArgumentParser
+from time import gmtime, strftime
+from shutil import copy
+from frames_dataset import FramesDataset
+
+from modules.inpainting_network import InpaintingNetwork
+from modules.keypoint_detector import KPDetector
+from modules.bg_motion_predictor import BGMotionPredictor
+from modules.dense_motion import DenseMotionNetwork
+from modules.avd_network import AVDNetwork
+import torch
+from train import train
+from train_avd import train_avd
+from reconstruction import reconstruction
+import os 
+
+
+if __name__ == "__main__":
+    
+    if sys.version_info[0] < 3:
+        raise Exception("You must use Python 3 or higher. Recommended version is Python 3.9")
+
+    parser = ArgumentParser()
+    parser.add_argument("--config", default="config/vox-256.yaml", help="path to config")
+    parser.add_argument("--mode", default="train", choices=["train", "reconstruction", "train_avd"])
+    parser.add_argument("--log_dir", default='log', help="path to log into")
+    parser.add_argument("--checkpoint", default=None, help="path to checkpoint to restore")
+    parser.add_argument("--device_ids", default="0,1", type=lambda x: list(map(int, x.split(','))),
+                        help="Names of the devices comma separated.")
+
+    opt = parser.parse_args()
+    with open(opt.config) as f:
+        config = yaml.load(f)
+
+    if opt.checkpoint is not None:
+        log_dir = os.path.join(*os.path.split(opt.checkpoint)[:-1])
+    else:
+        log_dir = os.path.join(opt.log_dir, os.path.basename(opt.config).split('.')[0])
+        log_dir += ' ' + strftime("%d_%m_%y_%H.%M.%S", gmtime())
+
+    inpainting = InpaintingNetwork(**config['model_params']['generator_params'],
+                                        **config['model_params']['common_params'])
+
+    if torch.cuda.is_available():
+        cuda_device = torch.device('cuda:'+str(opt.device_ids[0]))
+        inpainting.to(cuda_device)
+
+    kp_detector = KPDetector(**config['model_params']['common_params'])
+    dense_motion_network = DenseMotionNetwork(**config['model_params']['common_params'],
+                                              **config['model_params']['dense_motion_params'])
+                                                           
+    if torch.cuda.is_available():
+        kp_detector.to(opt.device_ids[0])
+        dense_motion_network.to(opt.device_ids[0])
+
+    bg_predictor = None
+    if (config['model_params']['common_params']['bg']):
+        bg_predictor = BGMotionPredictor()
+        if torch.cuda.is_available():
+            bg_predictor.to(opt.device_ids[0])
+
+    avd_network = None
+    if opt.mode == "train_avd":
+        avd_network = AVDNetwork(num_tps=config['model_params']['common_params']['num_tps'],
+                             **config['model_params']['avd_network_params'])
+        if torch.cuda.is_available():
+            avd_network.to(opt.device_ids[0])
+
+    dataset = FramesDataset(is_train=(opt.mode.startswith('train')), **config['dataset_params'])
+
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+    if not os.path.exists(os.path.join(log_dir, os.path.basename(opt.config))):
+        copy(opt.config, log_dir)
+
+    if opt.mode == 'train':
+        print("Training...")
+        train(config, inpainting, kp_detector, bg_predictor, dense_motion_network, opt.checkpoint, log_dir, dataset)
+    elif opt.mode == 'train_avd':
+        print("Training Animation via Disentaglement...")
+        train_avd(config, inpainting, kp_detector, bg_predictor, dense_motion_network, avd_network, opt.checkpoint, log_dir, dataset)
+    elif opt.mode == 'reconstruction':
+        print("Reconstruction...")
+        reconstruction(config, inpainting, kp_detector, bg_predictor, dense_motion_network, opt.checkpoint, log_dir, dataset)

train.py

@@ -0,0 +1,94 @@
+from tqdm import trange
+import torch
+from torch.utils.data import DataLoader
+from logger import Logger
+from modules.model import GeneratorFullModel
+from torch.optim.lr_scheduler import MultiStepLR
+from torch.nn.utils import clip_grad_norm_
+from frames_dataset import DatasetRepeater
+import math
+
+def train(config, inpainting_network, kp_detector, bg_predictor, dense_motion_network, checkpoint, log_dir, dataset):
+    train_params = config['train_params']
+    optimizer = torch.optim.Adam(
+        [{'params': list(inpainting_network.parameters()) +
+                    list(dense_motion_network.parameters()) +
+                    list(kp_detector.parameters()), 'initial_lr': train_params['lr_generator']}],lr=train_params['lr_generator'], betas=(0.5, 0.999), weight_decay = 1e-4)
+    
+    optimizer_bg_predictor = None
+    if bg_predictor:
+        optimizer_bg_predictor = torch.optim.Adam(
+            [{'params':bg_predictor.parameters(),'initial_lr': train_params['lr_generator']}], 
+            lr=train_params['lr_generator'], betas=(0.5, 0.999), weight_decay = 1e-4)
+
+    if checkpoint is not None:
+        start_epoch = Logger.load_cpk(
+            checkpoint, inpainting_network = inpainting_network, dense_motion_network = dense_motion_network,       
+            kp_detector = kp_detector, bg_predictor = bg_predictor,
+            optimizer = optimizer, optimizer_bg_predictor = optimizer_bg_predictor)
+        print('load success:', start_epoch)
+        start_epoch += 1
+    else:
+        start_epoch = 0
+
+    scheduler_optimizer = MultiStepLR(optimizer, train_params['epoch_milestones'], gamma=0.1,
+                                      last_epoch=start_epoch - 1)
+    if bg_predictor:
+        scheduler_bg_predictor = MultiStepLR(optimizer_bg_predictor, train_params['epoch_milestones'],
+                                              gamma=0.1, last_epoch=start_epoch - 1)
+
+    if 'num_repeats' in train_params or train_params['num_repeats'] != 1:
+        dataset = DatasetRepeater(dataset, train_params['num_repeats'])
+    dataloader = DataLoader(dataset, batch_size=train_params['batch_size'], shuffle=True, 
+                            num_workers=train_params['dataloader_workers'], drop_last=True)
+
+    generator_full = GeneratorFullModel(kp_detector, bg_predictor, dense_motion_network, inpainting_network, train_params)
+
+    if torch.cuda.is_available():
+        generator_full = torch.nn.DataParallel(generator_full).cuda()  
+        
+    bg_start = train_params['bg_start']
+    
+    with Logger(log_dir=log_dir, visualizer_params=config['visualizer_params'], 
+                checkpoint_freq=train_params['checkpoint_freq']) as logger:
+        for epoch in trange(start_epoch, train_params['num_epochs']):
+            for x in dataloader:
+                if(torch.cuda.is_available()):
+                    x['driving'] = x['driving'].cuda()
+                    x['source'] = x['source'].cuda()
+
+                losses_generator, generated = generator_full(x, epoch)
+                loss_values = [val.mean() for val in losses_generator.values()]
+                loss = sum(loss_values)
+                loss.backward()
+
+                clip_grad_norm_(kp_detector.parameters(), max_norm=10, norm_type = math.inf)
+                clip_grad_norm_(dense_motion_network.parameters(), max_norm=10, norm_type = math.inf)
+                if bg_predictor and epoch>=bg_start:
+                    clip_grad_norm_(bg_predictor.parameters(), max_norm=10, norm_type = math.inf)
+                
+                optimizer.step()
+                optimizer.zero_grad()
+                if bg_predictor and epoch>=bg_start:
+                    optimizer_bg_predictor.step()
+                    optimizer_bg_predictor.zero_grad()
+                
+                losses = {key: value.mean().detach().data.cpu().numpy() for key, value in losses_generator.items()}
+                logger.log_iter(losses=losses)
+
+            scheduler_optimizer.step()
+            if bg_predictor:
+                scheduler_bg_predictor.step()
+            
+            model_save = {
+                'inpainting_network': inpainting_network,
+                'dense_motion_network': dense_motion_network,
+                'kp_detector': kp_detector,
+                'optimizer': optimizer,
+            }
+            if bg_predictor and epoch>=bg_start:
+                model_save['bg_predictor'] = bg_predictor
+                model_save['optimizer_bg_predictor'] = optimizer_bg_predictor
+            
+            logger.log_epoch(epoch, model_save, inp=x, out=generated)
+

train_avd.py

@@ -0,0 +1,91 @@
+from tqdm import trange
+import torch
+from torch.utils.data import DataLoader
+from logger import Logger
+from torch.optim.lr_scheduler import MultiStepLR
+from frames_dataset import DatasetRepeater
+
+
+def random_scale(kp_params, scale):
+    theta = torch.rand(kp_params['fg_kp'].shape[0], 2) * (2 * scale) + (1 - scale)
+    theta = torch.diag_embed(theta).unsqueeze(1).type(kp_params['fg_kp'].type())
+    new_kp_params = {'fg_kp': torch.matmul(theta, kp_params['fg_kp'].unsqueeze(-1)).squeeze(-1)}
+    return new_kp_params
+
+
+def train_avd(config, inpainting_network, kp_detector, bg_predictor, dense_motion_network, 
+              avd_network, checkpoint, log_dir, dataset):
+    train_params = config['train_avd_params']
+
+    optimizer = torch.optim.Adam(avd_network.parameters(), lr=train_params['lr'], betas=(0.5, 0.999))
+
+    if checkpoint is not None:
+        Logger.load_cpk(checkpoint, inpainting_network=inpainting_network, kp_detector=kp_detector,
+                                      bg_predictor=bg_predictor, avd_network=avd_network,
+                                      dense_motion_network= dense_motion_network,optimizer_avd=optimizer)
+        start_epoch = 0
+    else:
+        raise AttributeError("Checkpoint should be specified for mode='train_avd'.")
+
+    scheduler = MultiStepLR(optimizer, train_params['epoch_milestones'], gamma=0.1)
+
+    if 'num_repeats' in train_params or train_params['num_repeats'] != 1:
+        dataset = DatasetRepeater(dataset, train_params['num_repeats'])
+
+    dataloader = DataLoader(dataset, batch_size=train_params['batch_size'], shuffle=True,
+                            num_workers=train_params['dataloader_workers'], drop_last=True)
+
+    with Logger(log_dir=log_dir, visualizer_params=config['visualizer_params'], 
+                checkpoint_freq=train_params['checkpoint_freq']) as logger:
+        for epoch in trange(start_epoch, train_params['num_epochs']):
+            avd_network.train()
+            for x in dataloader:
+                with torch.no_grad():
+                    kp_source = kp_detector(x['source'].cuda())
+                    kp_driving_gt = kp_detector(x['driving'].cuda())
+                    kp_driving_random = random_scale(kp_driving_gt, scale=train_params['random_scale'])
+                rec = avd_network(kp_source, kp_driving_random)
+
+                reconstruction_kp = train_params['lambda_shift'] * \
+                                       torch.abs(kp_driving_gt['fg_kp'] - rec['fg_kp']).mean()
+                
+                loss_dict = {'rec_kp': reconstruction_kp}
+                loss = reconstruction_kp
+                
+                loss.backward()
+                optimizer.step()
+                optimizer.zero_grad()
+
+                losses = {key: value.mean().detach().data.cpu().numpy() for key, value in loss_dict.items()}
+                logger.log_iter(losses=losses)
+
+            # Visualization
+            avd_network.eval()
+            with torch.no_grad():
+                source = x['source'][:6].cuda()
+                driving = torch.cat([x['driving'][[0, 1]].cuda(), source[[2, 3, 2, 1]]], dim=0)
+                kp_source = kp_detector(source)
+                kp_driving = kp_detector(driving)
+
+                out = avd_network(kp_source, kp_driving)
+                kp_driving = out
+                dense_motion = dense_motion_network(source_image=source, kp_driving=kp_driving,
+                                            kp_source=kp_source)
+                generated = inpainting_network(source, dense_motion)
+
+                generated.update({'kp_source': kp_source, 'kp_driving': kp_driving})
+
+            scheduler.step(epoch)
+            model_save = {
+                'inpainting_network': inpainting_network,
+                'dense_motion_network': dense_motion_network,
+                'kp_detector': kp_detector,
+                'avd_network': avd_network,
+                'optimizer_avd': optimizer
+            }
+            if bg_predictor :
+                model_save['bg_predictor'] = bg_predictor
+
+            logger.log_epoch(epoch, model_save,
+                             inp={'source': source, 'driving': driving},
+                             out=generated)