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") 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)