scripts/generate_video.py

from easing_functions.easing import LinearInOut
import torch
import pandas as pd 
from torchvision import utils as vutils
import os
import cv2
from tqdm import tqdm
from scipy import io
import numpy as np
import argparse

from easing_functions import QuadEaseInOut
from easing_functions import SineEaseIn, SineEaseInOut, SineEaseOut
from easing_functions import ElasticEaseIn, ElasticEaseInOut, ElasticEaseOut

ease_fn_dict = {'QuadEaseInOut': QuadEaseInOut,
                'SineEaseIn': SineEaseIn, 
                'SineEaseInOut': SineEaseInOut, 
                'SineEaseOut': SineEaseOut,
                'ElasticEaseIn': ElasticEaseIn, 
                'ElasticEaseInOut': ElasticEaseInOut, 
                'ElasticEaseOut': ElasticEaseOut,
                'Linear': LinearInOut}

def interpolate(z1, z2, num_interp):
    # this is a "first frame included, last frame excluded" interpolation
    w = torch.linspace(0, 1, num_interp+1)
    interp_zs = []
    for n in range(num_interp):
        interp_zs.append( (z2*w[n].item() + z1*(1-w[n].item())).unsqueeze(0) )
    return torch.cat(interp_zs)



def interpolate_ease_inout(z1, z2, num_interp, ease_fn, model_type='freeform'):
    # this is a "first frame included, last frame excluded" interpolation
    w = ease_fn(start=0, end=1, duration=num_interp+1)
    interp_zs = []

    # just to make sure the latent vectors in the right shape
    if model_type == 'freeform':
        z1 = z1.view(1, -1)
        z2 = z2.view(1, -1)
    if model_type == 'stylegan2':
        if type(z1) is list:
            z1 = [z1[0].view(1, -1), z1[1].view(1, -1)]
        else:
            z1 = [z1.view(1, -1), z1.view(1, -1)]
        if type(z2) is list:
            z2 = [z2[0].view(1, -1), z2[1].view(1, -1)]
        else:
            z2 = [z2.view(1, -1), z2.view(1, -1)]

    for n in range(num_interp):
        if model_type == 'freeform':
            interp_zs.append( z2*w.ease(n) + z1*(1-w.ease(n)) )
        if model_type == 'stylegan2':
            interp_zs.append( [ z2[0]*w.ease(n) + z1[0]*(1-w.ease(n)),
                                z2[1]*w.ease(n) + z1[1]*(1-w.ease(n)) ] )
    return interp_zs

@torch.no_grad()
def net_generate(netG, z, model_type='freeform', im_size=1024):
    
    if model_type == 'stylegan2':
        z_contents = []
        z_styles = []
        for zidx in range(len(z)):
            z_contents.append(z[zidx][0])
            z_styles.append(z[zidx][1])
        z = [ torch.cat(z_contents), torch.cat(z_styles) ]
        gimg = netG( z, inject_index=8, input_is_latent=True, randomize_noise=False )[0].cpu()
    elif model_type == 'freeform':
        z = torch.cat(z)
        gimg = netG(z)[0].cpu()

    return torch.nn.functional.interpolate(gimg, im_size)

def batch_generate_and_save(netG, zs, folder_name, batch_size=8, model_type='freeform', im_size=1024):
    # zs is a list of vectors if model is freeform
    # zs is a list of lists, each list is 2 vectors, if model is stylegan
    t = 0
    num = 0
    if len(zs) < batch_size:
        gimgs = net_generate(netG, zs, model_type, im_size=im_size).cpu()
        for image in gimgs:
            vutils.save_image( image.add(1).mul(0.5), folder_name+"/%d.jpg"%(num) )
            num += 1

    for k in tqdm(range(len(zs)//batch_size)):
        gimgs = net_generate(netG, zs[k*batch_size:(k+1)*batch_size], model_type, im_size=im_size)
        for image in gimgs:
            vutils.save_image( image.add(1).mul(0.5), folder_name+"/%d.jpg"%(num) )
            num += 1
        t = k

    if len(zs)%batch_size>0:
        gimgs = net_generate(netG, zs[(t+1)*batch_size:], model_type, im_size=im_size)        
        for image in gimgs:
            vutils.save_image( image.add(1).mul(0.5), folder_name+"/%d.jpg"%(num) )
            num += 1



def batch_save(images, folder_name, start_num=0):
    os.makedirs(folder_name, exist_ok=True)
    num = start_num
    for image in images:
        vutils.save_image( image.add(1).mul(0.5), folder_name+"/%d.jpg"%(num) )
        num += 1


def read_img_and_make_video(dist, video_name, fps):
    img_array = []
    for i in tqdm(range(len(os.listdir(dist)))):
        try:
            filename = dist+'/%d.jpg'%(i)
            img = cv2.imread(filename)
            height, width, layers = img.shape
            size = (width,height)
            img_array.append(img)
        except:
            print('error at: %d'%i)
    
    if '.mp4' not in video_name:
        video_name += '.mp4'
    out = cv2.VideoWriter(video_name,cv2.VideoWriter_fourcc(*'mp4v'), fps, size)
    for i in range(len(img_array)):
        out.write(img_array[i])
    out.release()

from shutil import rmtree

def make_video_from_latents(net, selected_latents, frames_dist_folder, video_name, fps, video_length, ease_fn, model_type, im_size=1024):
    # selected_latents: the latent noise of user selected key-frame images, it is a list
    # each item in the list is a vector if the model is freeform, 
    # each item in the list is a list of two vectors if the model is stylegan2
    # frames_dist_folder: the folder path to save the generated images to make the video
    # fps: is the frames we generate per second
    # video_length: is the time of the video, in seconds. For example: 30 means a video length of 30 seconds
    # ease_fn: user selected type of transitions between each key-frame

    # first calculate how many images need to generate
    try:
        rmtree(frames_dist_folder)
    except:
        pass
    os.makedirs(frames_dist_folder, exist_ok=True)

    nbr_generate = fps*video_length
    nbr_keyframe = len(selected_latents)
    nbr_interpolation = 1 + nbr_generate // (nbr_keyframe - 1)
    

    main_zs = []
    for idx in range(nbr_keyframe-1):
        main_zs += interpolate_ease_inout(selected_latents[idx], 
                            selected_latents[idx+1], nbr_interpolation, ease_fn, model_type)
    

    print('generating images ...')
    batch_generate_and_save(net, main_zs, folder_name=frames_dist_folder, batch_size=8, model_type=model_type, im_size=im_size)
    print('making videos ...')
    read_img_and_make_video(frames_dist_folder, video_name, fps=fps)


if __name__ == "__main__":


    device = torch.device('cuda:%d'%(0))

    load_model_err = 0

    from models import Generator as Generator_freeform
    
    frames_dist_folder = 'project_video_frames' # a folder to save generated images
    ckpt_path = './time_1024_1/models/180000.pth' # path to the checkpoint
    video_name = 'videl_keyframe_15'  # name of the generated video

    model_type = 'freeform'
    net = Generator_freeform(ngf=64, nz=100)
    net.load_state_dict(torch.load(ckpt_path)['g'])
    net.to(device)
    net.eval()

    
    try:
        rmtree(frames_dist_folder)
    except:
        pass
    os.makedirs(frames_dist_folder, exist_ok=True)

    fps = 30
    minutes = 1
    im_size = 1024
    
    ease_fn=ease_fn_dict['SineEaseInOut']

    init_kf_nbr = 15
    nbr_key_frames_per_minute = [init_kf_nbr-i for i in range(minutes)]
    nbr_key_frames_total = sum(nbr_key_frames_per_minute)
    noises = torch.randn( nbr_key_frames_total , 100).to(device)
    user_selected_noises = [n for n in noises]
    nbr_interpolation_list = [[fps*60//nbr_kf]*nbr_kf for nbr_kf in nbr_key_frames_per_minute]
    nbl = []
    for nb in nbr_interpolation_list:
        nbl += nb

    print(len(nbl)) 
    print(len(user_selected_noises))# , print("mismatch size")
    main_zs = []
    for idx in range(len(user_selected_noises)-1):
        main_zs += interpolate_ease_inout(user_selected_noises[idx], 
                            user_selected_noises[idx+1], nbl[idx], ease_fn, model_type)
    for idx in range(100):
        main_zs.append(main_zs[-1])
    print('generating images ...')
    batch_generate_and_save(net, main_zs, folder_name=frames_dist_folder, batch_size=8, model_type=model_type, im_size=im_size)
    print('making videos ...')
    read_img_and_make_video(frames_dist_folder, video_name, fps=fps)