text: # -*- coding: utf-8 -*-
  """LipGAN_Training.ipynb
  
  Automatically generated by Colaboratory.
  
  Original file is located at
      https://colab.research.google.com/drive/1KfhzyqZA6SNVGBcy1VXy2YK3aex2sss2
  """
  
  ### CONSTANTS
  filenames_filename = 'filenames.pkl'
  checkpoint_name = 'residual_model.h5'
  logs_dir = 'drive/MyDrive/BTP21SRD01/LipGAN/Training/logs/'
  resume_model = logs_dir+checkpoint_name
  checkpoint_freq = 1000
  
  img_size = 96
  half_window_size = 4
  mel_step_size = 27
  batch_size = 96
  lr = 1e-3
  img_size = 96
  epochs = 20
  
  """### Import Data"""
  
  !rm -Rf sample_data
  
  ### PATHS
  preprocessed_path = 'drive/MyDrive/BTP21SRD01/LipGAN/Preprocessing/*.zip'
  filenames_path = 'drive/MyDrive/BTP21SRD01/LipGAN/Preprocessing/' + filenames_filename
  
  !cp $filenames_path .
  !cp $preprocessed_path .
  
  !unzip -q preprocessed1.zip
  !unzip -q preprocessed2.zip
  !unzip -q preprocessed3.zip
  !unzip -q preprocessed4.zip
  !unzip -q preprocessed5.zip
  !rm *.zip
  
  """### Functions to Create Model"""
  
  import os
  import numpy as np
  import cv2
  import librosa
  import scipy
  import tensorflow as tf
  from tensorflow.keras.models import load_model
  from tensorflow.keras.optimizers import Adam
  from tensorflow.keras.models import Model
  from tensorflow.keras.layers import Dense, Conv2DTranspose, Conv2D, BatchNormalization,\
                          Activation, Concatenate, Input, MaxPool2D,\
  						UpSampling2D, ZeroPadding2D, Lambda, Add
  from tensorflow.keras.callbacks import ModelCheckpoint, Callback
  from tensorflow.keras import backend as K
  from tensorflow.keras.utils import plot_model
  
  def conv_block(x, num_filters, kernel_size=3, strides=1, padding='same', act=True):
      """
      function to add conv blocks
      Args:
          x -> tensor/ndarry: feature vector
          num_filters -> int: number of filters to be applied
          kernel_size -> int: filter size = (kernel_size x kernel_size) (default: 3)
          strides -> int: default 1
          padding-> str: 'same' to generate output with dimensions same as input
          act -> boolean: whether to use activation function or not
      """
  
      x = Conv2D(filters=num_filters, kernel_size= kernel_size, strides=strides, padding=padding)(x)
      x = BatchNormalization(momentum=.8)(x)
      if act:
          x = Activation('relu')(x)
      return x
  
  def conv_t_block(x, num_filters, kernel_size=3, strides=2, padding='same'):
      """
      function to add conv transpose blocks
      Args:
          x -> tensor/ndarry: feature vector
          num_filters -> int: number of filters to be applied
          kernel_size -> int: filter size = (kernel_size x kernel_size) (default: 3)
          strides -> int: default 1
          padding-> str: 'same' to generate output with dimensions same as input
          act -> boolean: whether to use activation function or not
      """
  
      x = Conv2DTranspose(filters=num_filters, kernel_size= kernel_size, strides=strides, padding=padding)(x)
      x = BatchNormalization(momentum=.8)(x)
      x = Activation('relu')(x)
      return x
  
  def residual_block(inp, num_filters):
      x = conv_block(inp, num_filters)
      x = conv_block(x, num_filters)
      x = Add()([x, inp])
      x = Activation('relu') (x)
      return x
  
  def create_model(img_size=img_size, mel_step_size=mel_step_size):
      """
      function to create the model
      Parts:
          face encoder, audio encoder, decoder
      Inputs: face, audio
      Output: face
      Args:
          img_size -> input for face encoder: (img_size x img_size)
          mel_step_size -> input for audio encoder: (80, mel_step_size)
      """
  
      ### face encoder
      input_face = Input(shape=(img_size, img_size, 6), name="input_face")
  
      identity_mapping = conv_block(input_face, 32, kernel_size=7) # 96x96
  
      x1_face = conv_block(identity_mapping, 64, kernel_size=5, strides=2) # 48x48
      x1_face = residual_block(x1_face, 64)
      x1_face = residual_block(x1_face, 64)
  
      x2_face = conv_block(x1_face, 128, 3, 2) # 24x24
      x2_face = residual_block(x2_face, 128)
      x2_face = residual_block(x2_face, 128)
      x2_face = residual_block(x2_face, 128)
  
      x3_face = conv_block(x2_face, 256, 3, 2) #12x12
      x3_face = residual_block(x3_face, 256)
      x3_face = residual_block(x3_face, 256)
  
      x4_face = conv_block(x3_face, 512, 3, 2) #6x6
      x4_face = residual_block(x4_face, 512)
      x4_face = residual_block(x4_face, 512)
  
      x5_face = conv_block(x4_face, 512, 3, 2) #3x3
      x6_face = conv_block(x5_face, 512, 3, 1, padding='valid')
      x7_face = conv_block(x6_face, 512, 1, 1)
  
      ### audio encoder
      input_audio = Input(shape=(80, mel_step_size, 1), name="input_audio")
  
      x = conv_block(input_audio, 32)
      x = residual_block(x, 32)
      x = residual_block(x, 32)
  
      x = conv_block(x, 64, strides=3)	#27X9
      x = residual_block(x, 64)
      x = residual_block(x, 64)
  
      x = conv_block(x, 128, strides=(3, 1)) 		#9X9
      x = residual_block(x, 128)
      x = residual_block(x, 128)
  
      x = conv_block(x, 256, strides=3)	#3X3
      x = residual_block(x, 256)
      x = residual_block(x, 256)
  
      x = conv_block(x, 512, strides=1, padding='valid')	#1X1
      x = conv_block(x, 512, 1, 1)
  
      embedding = Concatenate(axis=3)([x7_face, x])   # 1024x1
  
      ### deocder
      x = conv_t_block(embedding, 512, 3, 3)# 3x3
      x = Concatenate(axis=3) ([x5_face, x]) 
  
      x = conv_t_block(x, 512) #6x6
      x = residual_block(x, 512)
      x = residual_block(x, 512)
      x = Concatenate(axis=3) ([x4_face, x])
  
      x = conv_t_block(x, 256) #12x12
      x = residual_block(x, 256)
      x = residual_block(x, 256)
      x = Concatenate(axis=3) ([x3_face, x])
  
      x = conv_t_block(x, 128) #24x24
      x = residual_block(x, 128)
      x = residual_block(x, 128)
      x = Concatenate(axis=3) ([x2_face, x])
  
      x = conv_t_block(x, 64) #48x48
      x = residual_block(x, 64)
      x = residual_block(x, 64)
      x = Concatenate(axis=3) ([x1_face, x])
  
      x = conv_t_block(x, 32) #96x96
      x = Concatenate(axis=3) ([identity_mapping, x])
      x = conv_block(x, 16) #96x96
      x = conv_block(x, 16) #96x96
  
      x = Conv2D(filters=3, kernel_size=1, strides=1, padding="same") (x)
      prediction = Activation("sigmoid", name="prediction")(x)    # 0/1
  
      model = Model(inputs=[input_face, input_audio], outputs=prediction)
      return model
  
  # test_model = create_model(96, 27)
  # test_model.summary()
  
  class WeightsSaver(Callback):
  	def __init__(self, N, weight_path):
  		self.N = N
  		self.batch = 0
  		self.weight_path = weight_path
  
  	def on_batch_end(self, batch, logs={}):
  		self.batch += 1
  		if self.batch % self.N == 0:
  			self.model.save_weights(self.weight_path)
  
  """### Functions to Load Data"""
  
  import os, sys
  import pickle
  from glob import glob
  import numpy as np
  import scipy
  import cv2
  from os import listdir, path
  
  # list of paths of all images
  all_images = pickle.load(open(filenames_filename, 'rb'))
  len(all_images), type(all_images)
  
  def frame_id(fname):
      """
      fname: x.jpg
      returns: int(x)
      """
      return int(os.path.basename(fname).split('.')[0])
  
  def choose_ip_frame(frames, gt_frame):
      """
      func: choose surrounding frames to the current frame
      frames: list of frames
      gt_frame: current frame
      returns: one of the choosen frames
      """
      selected_frames = [f for f in frames if np.abs(frame_id(gt_frame) - frame_id(f)) >= 6]
      return np.random.choice(selected_frames)
  
  def get_audio_segment(center_frame, spec):
      """
      func: get the audio segment corressponding to the current frame
      """
      center_frame_id = frame_id(center_frame)
      start_frame_id = center_frame_id - half_window_size
      start_idx = int((80./25.) * start_frame_id) # 25 is fps of LRS2
      end_idx = start_idx + mel_step_size
      return spec[:, start_idx : end_idx] if end_idx <= spec.shape[1] else None
  
  def datagen(all_images=all_images, batch_size=batch_size):
  	"""
      func: generate the data of the form ([input_batch,spec],target_batch)
      args:
          all_images -> list of all image paths
          batch_size -> int: batch size for training
      returns: object of type generator
  	"""
  	while(True):
  		np.random.shuffle(all_images)
  		batches = [all_images[i:i + batch_size] for i in range(0, len(all_images), batch_size)]     # make batches of bs=`batch_size`
  
  		for batch in batches:
  			img_gt_batch = []       # target frames
  			img_ip_batch = []       # generated frames
  			mel_batch = []       # corresponsing melspec
  			
              ### goal is to generate the target frame from similar frames
  			for img_name in batch:
  				gt_fname = os.path.basename(img_name)       # x.png
  				dir_name = img_name.replace(gt_fname, '')   # dir contraining x.png
  				frames = glob(dir_name + '/*.jpg')          # all `valid` frames of the video x.png belongs to
  				if len(frames) < 12:                        # skip videos having less than 12 valid frames
  					continue
  				
                  ### load and process the saved melspec
  				mel_fname = dir_name + '/mels.npz'
  				try:
  					mel = np.load(mel_fname)['spec']
  				except:
  					continue
  
  				mel = get_audio_segment(gt_fname, mel)     # get the spec corressponding to the current frame
  
  				if mel is None or mel.shape[1] != mel_step_size:
  					continue
  
  				if sum(np.isnan(mel.flatten())) > 0:
  					continue
  				
                  ### load original frame
  				img_gt = cv2.imread(img_name)
  				img_gt = cv2.resize(img_gt, (img_size, img_size))
  				
                  ### choose and load a frame
  				ip_fname = choose_ip_frame(frames, gt_fname)
  				img_ip = cv2.imread(ip_fname)
  				img_ip = cv2.resize(img_ip, (img_size, img_size))
  
  				### make batches
  				img_gt_batch.append(img_gt)
  				img_ip_batch.append(img_ip)
  				mel_batch.append(mel)
  
  			img_gt_batch = np.asarray(img_gt_batch)
  			img_ip_batch = np.asarray(img_ip_batch)
  			mel_batch = np.expand_dims(np.asarray(mel_batch), 3)
  
  			img_gt_batch_masked = img_gt_batch.copy()
  			img_gt_batch_masked[:, img_size//2:,...] = 0.       # add mask filters
  			img_ip_batch = np.concatenate([img_ip_batch, img_gt_batch_masked], axis=3)  # resultant images have 6 channels
  			
  			yield [img_ip_batch/255.0, mel_batch], img_gt_batch/255.0
  
  """### Training"""
  
  train_datagen = datagen()
  type(train_datagen)
  
  callbacks_list = [WeightsSaver(checkpoint_freq, os.path.join(logs_dir, checkpoint_name))]
  
  model = create_model()
  model.summary()
  
  model.compile(loss='mae', optimizer=Adam(learning_rate=lr))
  model.load_weights(resume_model)
  
  # model.fit_generator(
  #     train_datagen,
  #     steps_per_epoch=len(all_images)//batch_size,
  #     epochs=epochs,
  #     verbose=1,
  #     initial_epoch=0,
  #     callbacks = callbacks_list
  #     )
  
  model.fit_generator(
      train_datagen,
      steps_per_epoch=len(all_images)//batch_size,
      epochs=epochs,
      verbose=1,
      initial_epoch=9,
      callbacks = callbacks_list
      )
  
  