128 lines
3.9 KiB
Python
128 lines
3.9 KiB
Python
#encoding = utf8
|
|
import os
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
|
|
import face_detection
|
|
|
|
|
|
def read_files_path(path):
|
|
file_paths = []
|
|
files = os.listdir(path)
|
|
for file in files:
|
|
if not os.path.isdir(file):
|
|
file_paths.append(path + file)
|
|
return file_paths
|
|
|
|
|
|
def read_images(img_list):
|
|
frames = []
|
|
print('reading images...')
|
|
for img_path in tqdm(img_list):
|
|
print(f'read image path:{img_path}')
|
|
frame = cv2.imread(img_path)
|
|
frames.append(frame)
|
|
return frames
|
|
|
|
|
|
def mirror_index(size, index):
|
|
# size = len(self.coord_list_cycle)
|
|
turn = index // size
|
|
res = index % size
|
|
if turn % 2 == 0:
|
|
return res
|
|
else:
|
|
return size - res - 1
|
|
|
|
|
|
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, face_det_batch_size, pads):
|
|
detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
|
|
flip_input=False, device=device)
|
|
|
|
batch_size = face_det_batch_size
|
|
|
|
while 1:
|
|
predictions = []
|
|
try:
|
|
for i in tqdm(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 = pads
|
|
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)
|
|
if not False: 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_signal(frame, mel, face_det_results, img_size, wav2lip_batch_size=128):
|
|
img_batch, mel_batch, frame_batch, coord_batch = [], [], [], []
|
|
|
|
# for i, m in enumerate(mels):
|
|
idx = 0
|
|
frame_to_save = frame.copy()
|
|
face, coord = face_det_results[idx].copy()
|
|
|
|
face = cv2.resize(face, (img_size, img_size))
|
|
|
|
for i, m in enumerate(mel):
|
|
img_batch.append(face)
|
|
mel_batch.append(m)
|
|
frame_batch.append(frame_to_save)
|
|
coord_batch.append(coord)
|
|
|
|
if len(img_batch) >= wav2lip_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])
|
|
|
|
return img_batch, mel_batch, frame_batch, coord_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])
|
|
|
|
return img_batch, mel_batch, frame_batch, coord_batch
|
|
|