Spaces:
Sleeping
Sleeping
import os | |
import cv2 | |
from models import Wav2Lip | |
import torch | |
import audio | |
import numpy as np | |
import subprocess | |
import platform | |
import face_detection | |
def get_smoothened_boxes(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 face_detect(images, device): | |
detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, | |
flip_input=False, device=device) | |
batch_size = 16 | |
while 1: | |
predictions = [] | |
try: | |
for i in range(0, len(images), batch_size): | |
predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size]))) | |
except RuntimeError: | |
if batch_size == 1: | |
raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument') | |
batch_size //= 2 | |
print('Recovering from OOM error; New batch size: {}'.format(batch_size)) | |
continue | |
break | |
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) | |
results.append([x1, y1, x2, y2]) | |
boxes = np.array(results) | |
boxes = get_smoothened_boxes(boxes, T=5) | |
results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)] | |
del detector | |
return results | |
def datagen(frames, mels, batch_size, device): | |
img_size = 96 | |
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] | |
#face_det_results = face_detect([frames[0]], device) | |
face_det_results = face_detect(frames, device) | |
for i, m in enumerate(mels): | |
idx = i%len(frames) | |
frame_to_save = frames[idx].copy() | |
face, coords = face_det_results[idx].copy() | |
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) >= 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]) | |
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]) | |
yield img_batch, mel_batch, frame_batch, coords_batch | |
def _load(checkpoint_path, device): | |
if device == 'cuda': | |
checkpoint = torch.load(checkpoint_path) | |
else: | |
checkpoint = torch.load(checkpoint_path, | |
map_location=lambda storage, loc: storage) | |
return checkpoint | |
def load_lip_model(path="checkpoints/wav2lip_gan.pth", device="cuda"): | |
model = Wav2Lip() | |
print("Load checkpoint from: {}".format(path)) | |
checkpoint = _load(path, device) | |
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) | |
return model.eval() | |
def load_input_image_or_video(file_path, resize_factor=2, logging=False): | |
if not os.path.isfile(file_path): | |
print(file_path) | |
raise ValueError('input image or video must be a valid path to video/image file') | |
elif file_path.split('.')[1] in ['jpg', 'png', 'jpeg']: | |
full_frames = [cv2.imread(file_path)] | |
fps = 25 | |
else: | |
video_stream = cv2.VideoCapture(file_path) | |
fps = video_stream.get(cv2.CAP_PROP_FPS) | |
print('Reading video frames...') | |
full_frames = [] | |
while 1: | |
still_reading, frame = video_stream.read() | |
if not still_reading: | |
video_stream.release() | |
break | |
if resize_factor > 1: | |
frame = cv2.resize(frame, (frame.shape[1]//resize_factor, frame.shape[0]//resize_factor)) | |
full_frames.append(frame) | |
#print("Number of frames available for inference: "+str(len(full_frames))) | |
return full_frames, fps | |
def load_input_audio(file_path, fps, results_path): | |
mel_step_size = 16 | |
if not file_path.endswith('.wav'): | |
temp_file = f"{os.path.dirname(results_path)}/temp.wav" | |
command = 'ffmpeg -y -i {} -strict -2 {}'.format(file_path, temp_file) | |
subprocess.call(command, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) | |
file_path = temp_file | |
wav = audio.load_wav(file_path, sr=16000) | |
mel = audio.melspectrogram(wav) | |
if np.isnan(mel.reshape(-1)).sum() > 0: | |
raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again') | |
mel_chunks = [] | |
mel_idx_multiplier = 80./fps | |
i = 0 | |
while 1: | |
start_idx = int(i * mel_idx_multiplier) | |
if start_idx + mel_step_size > len(mel[0]): | |
mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:]) | |
break | |
mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size]) | |
i += 1 | |
#print("Length of mel chunks: {}".format(len(mel_chunks))) | |
return mel_chunks, file_path | |
def animate_input(frames, audio, audio_file, fps, model, device, results_path): | |
temp_path = f"{os.path.dirname(results_path)}" | |
full_frames = frames[:len(audio)] | |
batch_size = 128 | |
gen = datagen(full_frames.copy(), audio, batch_size, device) | |
for i, (img_batch, mel_batch, frames, coords) in enumerate(gen): | |
if i == 0: | |
temp_video = temp_path + "/result.avi" | |
frame_h, frame_w = full_frames[0].shape[:-1] | |
out = cv2.VideoWriter(temp_video, | |
cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h)) | |
img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device) | |
mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device) | |
with torch.no_grad(): | |
pred = model(mel_batch, img_batch) | |
pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255. | |
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)) | |
f[y1:y2, x1:x2] = p | |
out.write(f) | |
out.release() | |
command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(audio_file, temp_video, results_path) | |
subprocess.call(command, shell=platform.system() != 'Windows', stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) |