avatar.py

from gtts import gTTS
import cv2
from pydub import AudioSegment
from tqdm import tqdm
import numpy as np
import pickle
import time
import subprocess,platform
import os
import torch
import io
import soundfile as sf

from models import Wav2Lip
import face_detection
import audio


class Avatar:
  image_frame_num_current = 0
  image_frame_num_goal=0
  wav2lip_gan_model = []  # should be a model
  video_full_frames = []
  images_and_audio_list = []
  images_list = []
  mel_step_size = 16
  output_audio_path = ""
  output_audio_filename = ""
  temp_lip_video_no_voice_path=""
  temp_lip_video_no_voice_filename=""
  input_audio_path=""
  input_video_path=""
  output_video_path=""
  output_video_name=""
  lip_video_no_voice_path=""
  split_current_file_name = ""
  fps = 30.0
  face_detect_img_results = []
  device = ""
  face_detect_batch_size = 16
  face_det_results_path_and_name = ""
  datagen_batch_size = 512
  frame_count = 0
  video_width = 0
  video_height = 0
  export_video = False
  def __init__(self):
    print("Avatar init")

  def _load(self,checkpoint_path):
      device = 'cuda' if torch.cuda.is_available() else 'cpu'
      if device == 'cuda':
          checkpoint = torch.load(checkpoint_path)
      else:
          checkpoint = torch.load(checkpoint_path,
                                  map_location=lambda storage, loc: storage)
      return checkpoint


  def load_model(self,path):
      device = 'cuda' if torch.cuda.is_available() else 'cpu'
      model = Wav2Lip()
      print("Load checkpoint from: {}".format(path))
      checkpoint = self._load(path)
      s = checkpoint["state_dict"]
      new_s = {}
      for k, v in s.items():
          new_s[k.replace('module.', '')] = v
      model.load_state_dict(new_s)

      model = model.to(device)
      self.wav2lip_gan_model = model.eval()
    #return model.eval()

  def get_video_full_frames(self, video_path):
    video_stream = cv2.VideoCapture(video_path)

    self.fps = video_stream.get(cv2.CAP_PROP_FPS)
    self.frame_count = video_stream.get(cv2.CAP_PROP_FRAME_COUNT)
    self.video_width = video_stream.get(cv2.CAP_PROP_FRAME_WIDTH)
    self.video_height = video_stream.get(cv2.CAP_PROP_FRAME_HEIGHT)
    print("fps="+str(self.fps))
    print('Reading video frames...')

    self.video_full_frames = []

    is_first_frame = True

    while 1:
      still_reading, frame = video_stream.read()
      if not still_reading:
        video_stream.release()
        break
      self.video_full_frames.append(frame)
      # if is_first_frame:
      #   first_frame_shape=frame.shape
      #   first_frame=frame
      #   is_first_frame = False
    # #[] is a list
    # #full frame is a video!!!!, so 4 dimension
    # print("CV2 frame_count="+str(self.frame_count))
    # print("CV2 video_width="+str(self.video_width))
    # print("CV2 video_height="+str(self.video_height))
    # print("first frame shape="+str(first_frame_shape))
    #
    # print ("Number of frames available for inference: "+str(len(self.video_full_frames)))
    # print("frame element="+str(first_frame[100][100][0]))
    # print("frame element type="+str(type(first_frame[100][100][0])))
    # #The value range for numpy.uint8 is from 0 to 255.
    # print("len(full_frames)"+str(len(self.video_full_frames)))



  def create_mel_from_audio(self,input_text):
    tts = gTTS(text=input_text, lang="en")
    if not os.path.exists(self.output_audio_path):
      print(f"{self.output_audio_path} does not exist, creating one")
      os.makedirs(self.output_audio_path)
    tts.save(f"{self.output_audio_path}input_audio.mp3")
    sound = AudioSegment.from_mp3(f"{self.output_audio_path}input_audio.mp3")

    # Get the duration in seconds
    sound_duration = sound.duration_seconds

    sound.export(f"{self.output_audio_path}temp_{self.output_audio_filename}", format="wav")

    wav = audio.load_wav(f"{self.output_audio_path}temp_{self.output_audio_filename}", 16000)

    mel = audio.melspectrogram(wav)
    # print(mel.shape)
    # #(80, 97)
    # #It means that the mel spectrogram of the audio input has 80 mel frequency bands and 97 time frames.
    # #Yes, mel frequency bands do overlap
    # #if wav is longer, so will the nmber of time frames
    # #(80, 344)
    # #mel is numpy.ndarray
    # print("mel data type =" + str(type(mel)))
    # print("mel element type =" +str(type(mel[1][2])))
    # print("len(mel[0])="+str(len(mel[0])))
    # #each mel element is numpy.float64, so can go negative

    mel_chunks = []
    mel_idx_multiplier = 80./self.fps
    #seems there is always 80 mel frequency
    #print("mel_idx_multiplier="+str(mel_idx_multiplier))
    #30 frames per seconds, 80 mel frequency bands, so 2.66 bands per frame per second
    i = 0
    while 1:
      start_idx = int(i * mel_idx_multiplier)
      #len(mel[0]) is the number of time frames of the audio
      if start_idx + self.mel_step_size > len(mel[0]):
        mel_chunks.append(mel[:, len(mel[0]) - self.mel_step_size:])
        break
      mel_chunks.append(mel[:, start_idx : start_idx + self.mel_step_size])
      i += 1
    # for b_index, b_item in enumerate(reversed(mel_chunks)):
    #   print(str(b_index)+" "+str(np.average(b_item)))

    for index, item in enumerate(reversed(mel_chunks)):
      #print(str(index)+" "+str(np.average(item)))
      if np.average(item) > -4.0:
        break
    print("stop at "+str(index))
    num_frames_to_trim=index-1
    mel_chunks=mel_chunks[:-num_frames_to_trim]
    print("wav length={} duration={} num_frames_to_trim={} result={}".format(len(wav),sound_duration,num_frames_to_trim,str(16000*num_frames_to_trim//30)))
    wav=wav[:-(16000*num_frames_to_trim//30)]
    sf.write(f"{self.output_audio_path}{self.output_audio_filename}", wav, 16000)
    sound_file = io.BytesIO(open(f"{self.output_audio_path}{self.output_audio_filename}", "rb").read())

    # Load the wav file as an AudioSegment object
    audio_segment_sound = AudioSegment.from_wav(sound_file)
    return mel_chunks, audio_segment_sound

  def get_smoothened_boxes(self, boxes, T):
      for i in range(len(boxes)):
          if i + T > len(boxes):
              window = boxes[len(boxes) - T:]
          else:
              window = boxes[i : i + T]
          boxes[i] = np.mean(window, axis=0)
      return boxes


  def create_face_detection_results(self, full_frames,save_result=True):
    detector = FACE_DETECTION.FaceAlignment(FACE_DETECTION.LandmarksType._2D,
                                              flip_input=False, device=self.device)
    images=full_frames
    while 1:
      predictions = []
      try:
        for i in tqdm(range(0, len(images), self.face_detect_batch_size)):
          predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + self.face_detect_batch_size])))
      except RuntimeError:
        if self.face_detect_batch_size == 1:
          raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
        self.face_detect_batch_size //= 2
        print('Recovering from OOM error; New batch size: {}'.format(self.face_detect_batch_size))
        continue
      break

    face_detect_results = []
    pady1, pady2, padx1, padx2 = [0, 10, 0, 0]
    for rect, image in zip(predictions, images):
      if rect is None:
        cv2.imwrite('temp_faulty_frame.jpg', image) # check this frame where the face was not detected.
        raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')

      y1 = max(0, rect[1] - pady1)
      y2 = min(image.shape[0], rect[3] + pady2)
      x1 = max(0, rect[0] - padx1)
      x2 = min(image.shape[1], rect[2] + padx2)

      face_detect_results.append([x1, y1, x2, y2])
    # print("\n")
    # print("face_detect_results length = " + str(len(face_detect_results)))
    # print("face_detect_results[2]="+str(face_detect_results[2]))

    boxes = np.array(face_detect_results)
    boxes = self.get_smoothened_boxes(boxes, T=5)
    # print ("boxes number of dim="+str(boxes.ndim))
    # print ("boxes shape="+str(boxes.shape))

    self.face_detect_img_results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
    # print ("face_detect_img_results type =" + str(type(self.face_detect_img_results)))
    # print ("face_detect_img_results length =" + str(len(self.face_detect_img_results)))
    # print ("face_detect_img_results[1] type =" + str(type(self.face_detect_img_results[1])))
    # print ("face_detect_img_results[1] length =" + str(len(self.face_detect_img_results[1])))
    # print ("face_detect_img_results[1][1] = " +str(self.face_detect_img_results[1][1])) #this is the box
    # print ("face_detect_img_results[1][1] shape = " +str(self.face_detect_img_results[1][0].shape)) #this is cropped image
    if save_result:
      with open(self.face_det_results_path_and_name, 'wb') as file:
        pickle.dump(self.face_detect_img_results, file)


  def load_face_detection_results(self):
    with open(self.face_det_results_path_and_name, 'rb') as file:
      self.face_detect_img_results = pickle.load(file)
    # print ("face_detect_img_results type =" + str(type(self.face_detect_img_results)))
    # print ("face_detect_img_results length =" + str(len(self.face_detect_img_results)))
    # print ("face_detect_img_results[1] type =" + str(type(self.face_detect_img_results[1])))
    # print ("face_detect_img_results[1] length =" + str(len(self.face_detect_img_results[1])))
    # print ("face_detect_img_results[1][1] = " +str(self.face_detect_img_results[1][1])) #this is the box
    # print ("face_detect_img_results[1][1] shape = " +str(self.face_detect_img_results[1][0].shape)) #this is cropped image


  def datagen(self, full_frames, mels, face_detect_results):
    img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
    print(len(full_frames))
    for i, m in enumerate(mels):
      idx = i%len(full_frames)
      frame_to_save = full_frames[idx].copy()
      face, coords = face_detect_results[idx].copy()
      img_size = 96 # for wav2lip, their model is trained on 96x96 image
      face = cv2.resize(face, (img_size, img_size))

      img_batch.append(face)
      mel_batch.append(m)
      frame_batch.append(frame_to_save)
      coords_batch.append(coords)

      if len(img_batch) >= self.datagen_batch_size:

        img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)

        img_masked = img_batch.copy()
        img_masked[:, img_size//2:] = 0

        img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
        mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
        #print(f"len(img_batch)>{self.datagen_batch_size} now len(img_batch)=" + str(len(img_batch)))
        yield img_batch, mel_batch, frame_batch, coords_batch
        img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []

    if len(img_batch) > 0:
      img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)

      img_masked = img_batch.copy()
      img_masked[:, img_size//2:] = 0

      img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
      mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
      #print("len(img_batch)>0 now len(img_batch)="+str(len(img_batch)))
      yield img_batch, mel_batch, frame_batch, coords_batch
      # datagen_result=datagen(video_full_frames_copy, 512, mel_chunks_from_audio,face_results)
      # for img, mel, frame, coords in datagen_result:
      #   gen_img = img.copy()
      #   gen_mel = mel.copy()
      #   gen_frame = frame.copy()
      #   gen_coords = coords.copy()
      # print("gen image shape ="+str(gen_img.shape))
      # print("gen mel shape = "+str(gen_mel.shape))
      # print("gen_coords length = " + str(len(gen_coords)))
      # print("gen_coords[0] = " + str(gen_coords[0]))
      # print("gen_frame length =" + str(len(gen_frame)))
      # print("gen_frame[0] type =" + str(type(gen_frame[0])))
      # print("gen_frame[0] shape =" + str(gen_frame[0].shape))
      # print(str(gen_frame[0].shape))
      #You are seeing img_batch shape as (batch_size, 96, 96, 6) because you are using the Wav2Lip model with the face detection and alignment option enabled. This option preprocesses the face images by detecting the face region, aligning the face orientation, and cropping and resizing the face image to 96 by 96 pixels. However, instead of discarding the original face image, the option concatenates the aligned face image and the original face image along the channel dimension, resulting in a 6-channel image.



  def make_lip_video(self, datagen_result,video_write_out, mel_chunks,need_split, audio_sound):


    for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(datagen_result,
                            total=int(np.ceil(float(len(mel_chunks))/self.datagen_batch_size)))):
        #print("\nin the for loop to unpack datagen_result, only run once")


        img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(self.device)
        mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(self.device)
        inf_start_time = time.time()  # get the start time
        with torch.no_grad():
          pred = self.wav2lip_gan_model(mel_batch, img_batch)

        pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
        # print("type of pred"+str(type(pred)))
        # print("shape of pred"+str(pred.shape))
        for p, f, c in zip(pred, frames, coords):
          y1, y2, x1, x2 = c
          p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1)) #before the face was extracted in scaled down to 96x96
          f[y1:y2, x1:x2] = p #paste face back
          self.images_list.append(f)
          if need_split:
            self.image_frame_num_current = self.image_frame_num_current + 1
          if self.export_video:
            video_write_out.write(f)


        if need_split:
          #print("GLOBAL_IMAGE_FRAME_NUM_CURRENT=" + str(self.image_frame_num_current))
          if self.image_frame_num_current >= self.image_frame_num_goal:
            self.images_and_audio_list.append([self.images_list, audio_sound])
            self.images_list = []
            # print("video_write_out relase in need split")
            if self.export_video:
              video_write_out.release()
        else:
          self.images_and_audio_list.append([self.images_list, audio_sound])
          self.images_list = []
          #print("video_write_out relase")
          if self.export_video:
            video_write_out.release()
        inf_end_time = time.time()  # get the end time
        print(f"Inference time: {inf_end_time - inf_start_time} seconds")  # print the difference
        #print("img_batch length="+str(len(img_batch)))
  def sync_video_audio(self,input_audio_path_and_name, input_video_path_and_name, output_video_path_and_name):
    #ipdb.set_trace()
    command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(input_audio_path_and_name,input_video_path_and_name,
                                                                  output_video_path_and_name)
    subprocess.call(command, shell=platform.system() != 'Windows')
    # command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(input_audio_path, input_video_path,res[0]+'/temp'+res[1] )
    # subprocess.call(command, shell=platform.system() != 'Windows')
    # command = 'ffmpeg -ss 00:00:00 -t {} -i {} -c copy {}'.format(str(round(num_of_frames/fps,2)),res[0]+'/temp'+res[1],output_video_path )
    # subprocess.call(command, shell=platform.system() != 'Windows')


  def video_audio_adjust(self, mel_chunks_from_audio, frame_chunks_from_video):
    out_frame_chunks_from_video=frame_chunks_from_video.copy()
    out_faces_from_detect_results=self.face_detect_img_results.copy()
    audio_duration= len(mel_chunks_from_audio)
    video_duration= len(frame_chunks_from_video)
    # print("mel length="+str(audio_duration))
    # print("frame length="+str(video_duration))
    if audio_duration != video_duration:
      if audio_duration>video_duration:
        differece=audio_duration-video_duration
        # calculate how many times video should be concat
        times=differece/video_duration
        # create a file with video name and concat the video using ffmpeg

        # if times fraction then add 1 to times
        if times%1!=0:
            times+=2
        out_frame_chunks_from_video=out_frame_chunks_from_video*int(times)
        out_faces_from_detect_results=out_faces_from_detect_results*int(times)
        new_video_duration= len(out_frame_chunks_from_video)
        # print("extending video frames and face detect")
        # print("new frame length="+str(new_video_duration))
        if new_video_duration > audio_duration:
          out_frame_chunks_from_video=out_frame_chunks_from_video[:audio_duration]
          out_faces_from_detect_results=out_faces_from_detect_results[:audio_duration]
        new_video_duration= len(out_frame_chunks_from_video)
        # print("new frame length="+str(new_video_duration))
      else:
        # print("truncate video frames and face detect")
        out_frame_chunks_from_video=out_frame_chunks_from_video[:audio_duration]
        out_faces_from_detect_results=out_faces_from_detect_results[:audio_duration]
    return out_frame_chunks_from_video,out_faces_from_detect_results

  def video_audio_adjust_parallel(self, mel_chunks_from_audio, frame_chunks_from_video,
                                  pre_frame_chunks_from_video, pre_faces_from_detect_results):

    out_frame_chunks_from_video=frame_chunks_from_video.copy()
    out_faces_from_detect_results=self.face_detect_img_results.copy()
    audio_duration= len(mel_chunks_from_audio)
    video_duration= len(frame_chunks_from_video)

    post_frame_chunks_from_video=[]
    post_faces_from_detect_results=[]
    pre_video_duration= len(pre_frame_chunks_from_video)

    # print("video_audio_adjust_parallel mel length="+str(audio_duration))
    # print("video_audio_adjust_parallel frame length="+str(video_duration))
    # print("video_audio_adjust_parallel pre frame length="+str(pre_video_duration))
    if (audio_duration-pre_video_duration) != video_duration:
      if (audio_duration-pre_video_duration)>video_duration:
        # print("in Case 1")
        differece=(audio_duration-pre_video_duration)-video_duration
        # calculate how many times video should be concat
        times=differece/video_duration
        # create a file with video name and concat the video using ffmpeg

        # if times fraction then add 1 to times
        if times%1!=0:
            times+=2
        # print("video_audio_adjust_parallel times="+str(times))
        out_frame_chunks_from_video=out_frame_chunks_from_video*int(times)
        # print("video_audio_adjust_parallel video length after multiplying with time ="+str(len(out_frame_chunks_from_video)))
        # if len(pre_frame_chunks_from_video) > 0 :
        #   cv2.imwrite('/content/pre_video_first_'+DEBUG_GLOBAL_CURRENT_FILE_NAME, pre_frame_chunks_from_video[0])
        #   cv2.imwrite('/content/pre_video_last_'+DEBUG_GLOBAL_CURRENT_FILE_NAME, pre_frame_chunks_from_video[-1])
        # cv2.imwrite('/content/video_first_'+DEBUG_GLOBAL_CURRENT_FILE_NAME, out_frame_chunks_from_video[0])
        out_frame_chunks_from_video=pre_frame_chunks_from_video+out_frame_chunks_from_video
        out_faces_from_detect_results=out_faces_from_detect_results*int(times)
        out_faces_from_detect_results=pre_faces_from_detect_results+out_faces_from_detect_results
        new_video_duration= len(out_frame_chunks_from_video)
        # print("extending video frames and face detect")
        # print("new frame length="+str(new_video_duration))
        if new_video_duration > audio_duration:
          # print("in Case 1a")
          c = np.absolute(out_frame_chunks_from_video[audio_duration-1]- out_frame_chunks_from_video[audio_duration]) # or c = a - b
          # print("video_audio_adjust_parallel difference at cut off is "+str(np.mean(c) ))

          out_frame_chunks_from_video_copy=out_frame_chunks_from_video.copy()
          out_frame_chunks_from_video=out_frame_chunks_from_video[:audio_duration]
          post_frame_chunks_from_video=out_frame_chunks_from_video_copy[audio_duration:]
          out_faces_from_detect_results_copy=out_faces_from_detect_results.copy()
          out_faces_from_detect_results=out_faces_from_detect_results[:audio_duration]
          post_faces_from_detect_results=out_faces_from_detect_results_copy[audio_duration:]

        #else:
          # print("unhandled case 1")
        #new_video_duration= len(out_frame_chunks_from_video)
        # print("new frame length="+str(new_video_duration))
      else:
        # print("in Case 2")
        # print("truncate video frames and face detect")

        # print("video_audio_adjust_parallel video length pre_frame_chunks_from_video ="+str(len(pre_frame_chunks_from_video)))
        # print("video_audio_adjust_parallel video length out_frame_chunks_from_video ="+str(len(out_frame_chunks_from_video)))

        out_frame_chunks_from_video=pre_frame_chunks_from_video+out_frame_chunks_from_video
        c = np.absolute(out_frame_chunks_from_video[audio_duration-1]-out_frame_chunks_from_video[audio_duration]) # or c = a - b
        # print("video_audio_adjust_parallel difference at cut off is "+str(np.mean(c) ))
        out_faces_from_detect_results=pre_faces_from_detect_results+out_faces_from_detect_results
        out_frame_chunks_from_video_copy=out_frame_chunks_from_video.copy()
        out_frame_chunks_from_video=out_frame_chunks_from_video[:audio_duration]
        post_frame_chunks_from_video=out_frame_chunks_from_video_copy[audio_duration:]




        out_faces_from_detect_results_copy=out_faces_from_detect_results.copy()
        out_faces_from_detect_results=out_faces_from_detect_results[:audio_duration]
        post_faces_from_detect_results=out_faces_from_detect_results_copy[audio_duration:]
        # cv2.imwrite('/content/video_last_'+DEBUG_GLOBAL_CURRENT_FILE_NAME, out_frame_chunks_from_video[-1])
        # if len(post_frame_chunks_from_video) > 0 :
        #   cv2.imwrite('/content/post_video_first_'+DEBUG_GLOBAL_CURRENT_FILE_NAME, post_frame_chunks_from_video[0])
        #   cv2.imwrite('/content/post_video_last_'+DEBUG_GLOBAL_CURRENT_FILE_NAME, post_frame_chunks_from_video[-1])
    #else:
      # print("unhandled case 2")
    return out_frame_chunks_from_video,out_faces_from_detect_results,post_frame_chunks_from_video,post_faces_from_detect_results


  def text_to_lip_video(self, input_text):

    mel_chunks_from_audio, audio_segment =self.create_mel_from_audio(input_text)
    print(str(len(self.face_detect_img_results)))
    video_full_frames_copy=self.video_full_frames.copy()

    video_full_frames_copy,face_detect_results=self.video_audio_adjust(mel_chunks_from_audio,video_full_frames_copy)

    gen=self.datagen(video_full_frames_copy, mel_chunks_from_audio,face_detect_results)

    if self.export_video:
      video_write_handle = cv2.VideoWriter(self.temp_lip_video_no_voice_path+self.temp_lip_video_no_voice_filename,
                            cv2.VideoWriter_fourcc(*'DIVX'), self.fps,
                                           (self.video_full_frames[0].shape[0], self.video_full_frames[0].shape[1]))
    else:
      video_write_handle =0

    self.make_lip_video(gen,video_write_handle, mel_chunks_from_audio,len(mel_chunks_from_audio)>self.datagen_batch_size,
                        audio_segment)
    if self.export_video:
      self.sync_video_audio(self.output_audio_path+self.output_audio_filename,
                            self.temp_lip_video_no_voice_path+self.temp_lip_video_no_voice_filename,
                            self.output_video_path+self.output_video_name
                            )
    self.image_frame_num_current=0
  def text_to_lip_video_parallel(self, input_text, pre_base_video_frames,pre_face_detect_results):



    mel_chunks_from_audio, audio_segment = self.create_mel_from_audio(input_text)
    print(str(len(self.face_detect_img_results)))
    video_full_frames_copy=self.video_full_frames.copy()

    video_full_frames_copy,face_detect_results,post_base_video_frames,post_face_detect_results=(
      self.video_audio_adjust_parallel(mel_chunks_from_audio,video_full_frames_copy,
                                       pre_base_video_frames,pre_face_detect_results))

    gen=self.datagen(video_full_frames_copy,mel_chunks_from_audio,face_detect_results)

    if self.export_video:
      video_write_handle = cv2.VideoWriter(self.temp_lip_video_no_voice_path + self.temp_lip_video_no_voice_filename,
                                           cv2.VideoWriter_fourcc(*'DIVX'), self.fps,
                                           (self.video_full_frames[0].shape[0], self.video_full_frames[0].shape[1]))
    else:
      video_write_handle=0

    self.make_lip_video(gen,video_write_handle, mel_chunks_from_audio,len(mel_chunks_from_audio)>self.datagen_batch_size,audio_segment)
    if self.export_video:
      self.sync_video_audio(self.output_audio_path + self.output_audio_filename,
                            self.temp_lip_video_no_voice_path + self.temp_lip_video_no_voice_filename,
                            self.output_video_path + self.split_current_file_name)
    image_frame_num_current=0


    return post_base_video_frames,post_face_detect_results
  def delete_files_in_path(self,dir_path):
    # Get the list of files in the directory
    files = os.listdir(dir_path)

    # Check if the list is not empty
    if files:
      # Loop through the files
      for file in files:
        # Join the file name with the directory path
        file_path = os.path.join(dir_path, file)
        # Check if the file is a regular file (not a directory or a link)
        if os.path.isfile(file_path):
        # Delete the file
          os.remove(file_path)

  def dir_clean_up(self):
    if os.path.isdir(self.output_audio_path):
      self.delete_files_in_path(self.output_audio_path)
    else:
      os.mkdir(self.output_audio_path)

    if (self.export_video):
      if os.path.isdir(self.temp_lip_video_no_voice_path):
        self.delete_files_in_path(self.temp_lip_video_no_voice_path)
      else:
        os.mkdir(self.temp_lip_video_no_voice_path)

      if os.path.isdir(self.output_video_path):
        self.delete_files_in_path(self.output_video_path)
      else:
        os.mkdir(self.output_video_path)