
#[Super SloMo]
##High Quality Estimation of Multiple Intermediate Frames for Video Interpolation

import argparse

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn

from model import superslomo_half as superslomo
from model.extraction import center, ends, MPRNet
from utils import meanshift, load_image

from copy import deepcopy
import os, glob
import numpy as np


# For parsing commandline arguments
parser = argparse.ArgumentParser()
parser.add_argument("--input", type=str, default='demo_input', help='path to dataset folder containing train-test-validation folders')
parser.add_argument("--output", type=str, default='demo_output', help='path to folder for saving checkpoints')
parser.add_argument("--checkpoint", type=str, help='path of checkpoint for pretrained model')
parser.add_argument("--add_blur", action='store_true', help='Add blurry image')
parser.add_argument("--seq_len", type=int, default=7, help='number of frames that composes a sequence.')

args = parser.parse_args()

### For visualizing loss and interpolated frames
###Initialize flow computation and arbitrary-time flow interpolation CNNs.

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
flowComp = superslomo.UNet(6, 4)
flowComp.to(device)
if args.add_blur:
    ArbTimeFlowIntrp = superslomo.UNet(23, 5)
else:
    ArbTimeFlowIntrp = superslomo.UNet(20, 5)
ArbTimeFlowIntrp.to(device)


### Load Pretrained Extraction Models

# center_estimation = center.Center()
center_estimation = MPRNet.MPRNet()
border_estimation = ends.Ends()
center_estimation = center_estimation.to(device)
border_estimation = border_estimation.to(device)

### Load Datasets
# Channel wise mean calculated on adobe240-fps training dataset
mean = [0.429, 0.431, 0.397]
std  = [1, 1, 1]

seq_len = args.seq_len
ctr_idx = seq_len // 2

def test(blurred_img, fname):
    with torch.no_grad():
        blurred_img = blurred_img.to(device)
        batch_size = blurred_img.shape[0]

        h, w = blurred_img.shape[-2:]

        ###Initialze backward warpers for train and validation datasets
        validationFlowBackWarp = superslomo.backWarp(w, h, device)
        validationFlowBackWarp = validationFlowBackWarp.to(device)

        blurred_img = meanshift(blurred_img, mean, std, device, False)
        center = center_estimation(blurred_img)
        start, end = border_estimation(blurred_img, center)

        start = meanshift(start, mean, std, device, True)
        end = meanshift(end, mean, std, device, True)
        center = meanshift(center, mean, std, device, True)
        blurred_img = meanshift(blurred_img, mean, std, device, True)
        
        I0 = torch.zeros_like(blurred_img)
        I1 = center

        for vindex in range(seq_len):            
            if vindex == 0:
                Ft_p = torch.zeros_like(blurred_img)
                for b in range(batch_size):
                    Ft_p[b] = start[b].clone()

            elif vindex == seq_len-1:
                Ft_p = torch.zeros_like(blurred_img)
                for b in range(batch_size):
                    Ft_p[b] = end[b].clone()

            elif vindex == ctr_idx:
                Ft_p = center.clone()

            else:
                validationIndex = torch.ones(batch_size) * (vindex - 1)
                validationIndex = validationIndex.long()
                if vindex > ctr_idx:
                    validationIndex = seq_len - 3 - validationIndex

                for b in range(batch_size):
                    if vindex < ctr_idx:
                        I0[b] = start[b]
                    else:
                        I0[b] = end[b]

                flowOut = flowComp(torch.cat((I0, I1), dim=1))
                F_0_1 = flowOut[:,:2,:,:]
                F_1_0 = flowOut[:,2:,:,:]

                fCoeff = superslomo.getFlowCoeff(validationIndex, device, seq_len)

                F_t_0 = fCoeff[0] * F_0_1 + fCoeff[1] * F_1_0
                F_t_1 = fCoeff[2] * F_0_1 + fCoeff[3] * F_1_0

                g_I0_F_t_0 = validationFlowBackWarp(I0, F_t_0)
                g_I1_F_t_1 = validationFlowBackWarp(I1, F_t_1)
                
                if args.add_blur:
                    intrpOut = ArbTimeFlowIntrp(torch.cat((I0, I1, F_0_1, F_1_0, F_t_1, F_t_0, g_I1_F_t_1, g_I0_F_t_0, blurred_img), dim=1))
                else:
                    intrpOut = ArbTimeFlowIntrp(torch.cat((I0, I1, F_0_1, F_1_0, F_t_1, F_t_0, g_I1_F_t_1, g_I0_F_t_0), dim=1))
                
                F_t_0_f = intrpOut[:, :2, :, :] + F_t_0
                F_t_1_f = intrpOut[:, 2:4, :, :] + F_t_1
                V_t_0   = torch.sigmoid(intrpOut[:, 4:5, :, :])
                V_t_1   = 1 - V_t_0
                
                g_I0_F_t_0_f = validationFlowBackWarp(I0, F_t_0_f)
                g_I1_F_t_1_f = validationFlowBackWarp(I1, F_t_1_f)
                
                wCoeff = superslomo.getWarpCoeff(validationIndex, device, seq_len)
                
                Ft_p = (wCoeff[0] * V_t_0 * g_I0_F_t_0_f + wCoeff[1] * V_t_1 * g_I1_F_t_1_f) / (wCoeff[0] * V_t_0 + wCoeff[1] * V_t_1)
            
            Ft_p = meanshift(Ft_p, mean, std, device, False)
            out_img = Ft_p.cpu().clone()

            filename = fname + '_est{:02d}'.format(vindex+1) + '.png'
            # Comment two lines below if you want to save images
            if not os.path.exists(os.path.join(args.output, fname)):
                os.mkdir(os.path.join(args.output, fname))
            torchvision.utils.save_image(out_img[0], os.path.join(args.output, fname, filename))


### Initialization
dict1 = torch.load(args.checkpoint)
ArbTimeFlowIntrp.load_state_dict(dict1['state_dictAT'])
flowComp.load_state_dict(dict1['state_dictFC'])
center_estimation.load_state_dict(dict1['state_dictCT'])
border_estimation.load_state_dict(dict1['state_dictBD'])
print('Load model from: ', args.checkpoint)

inputfolder = args.input
img_list = sorted(glob.glob(os.path.join(inputfolder, '*.png')))

args.output = os.path.join(args.output, 'Len{:02d}'.format(args.seq_len))
if not os.path.exists(args.output):
    os.makedirs(args.output)

for inputFile in img_list:
    print(inputFile)
    inputs = load_image(inputFile)
    fname = os.path.splitext(os.path.basename(inputFile))[0]
    width, height= inputs.size
    inputs = inputs.crop((0,0, width-width%64, height-height%64))

    normalize = transforms.Normalize(mean=mean, std=std)
    input_transform = transforms.Compose([transforms.ToTensor(), normalize])
    inputs = input_transform(inputs)
    inputs = inputs.unsqueeze(0)

    test(inputs, fname)
