bleckhert/Free-View_Expressive_Talking_Head_Video_Editing

1import face_detection2import numpy as np3import cv24from tqdm import tqdm5import torch6import glob7import os8from natsort import natsorted9 10device = "cuda" if torch.cuda.is_available() else "cpu"11 12 13def get_squre_coords(coords, image, size=None, last_size=None):14    y1, y2, x1, x2 = coords15    w, h = x2 - x1, y2 - y116    center = (x1 + w // 2, y1 + h // 2)17    if size is None:18        size = (w + h) // 219    if last_size is not None:20        size = (w + h) // 221        size = (size - last_size) // 5 + last_size22    x1, y1 = center[0] - size // 2, center[1] - size // 223    x2, y2 = x1 + size, y1 + size24    return size, [y1, y2, x1, x2]25 26 27def get_smoothened_boxes(boxes, T):28    for i in range(len(boxes)):29        if i + T > len(boxes):30            window = boxes[len(boxes) - T :]31        else:32            window = boxes[i : i + T]33        boxes[i] = np.mean(window, axis=0)34    return boxes35 36 37def face_detect(images, pads):38    detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, flip_input=False, device=device)39 40    batch_size = 32 if device == "cuda" else 441    print("face detect batch size:", batch_size)42    while 1:43        predictions = []44        try:45            for i in tqdm(range(0, len(images), batch_size)):46                predictions.extend(detector.get_detections_for_batch(np.array(images[i : i + batch_size])))47        except RuntimeError:48            if batch_size == 1:49                raise RuntimeError("Image too big to run face detection on GPU. Please use the --resize_factor argument")50            batch_size //= 251            print("Recovering from OOM error; New batch size: {}".format(batch_size))52            continue53        break54 55    results = []56    pady1, pady2, padx1, padx2 = pads57    for rect, image in zip(predictions, images):58        if rect is None:59            cv2.imwrite(".temp/faulty_frame.jpg", image)  # check this frame where the face was not detected.60            raise ValueError("Face not detected! Ensure the video contains a face in all the frames.")61 62        y1 = max(0, rect[1] - pady1)63        y2 = min(image.shape[0], rect[3] + pady2)64        x1 = max(0, rect[0] - padx1)65        x2 = min(image.shape[1], rect[2] + padx2)66        # y_gap, x_gap = ((y2 - y1) * 2) // 3, ((x2 - x1) * 2) // 367        y_gap, x_gap = (y2 - y1) // 2, (x2 - x1) // 268        coords_ = [y1 - y_gap, y2 + y_gap, x1 - x_gap, x2 + x_gap]69 70        _, coords = get_squre_coords(coords_, image)71 72        y1, y2, x1, x2 = coords73        y1 = max(0, y1)74        y2 = min(image.shape[0], y2)75        x1 = max(0, x1)76        x2 = min(image.shape[1], x2)77 78        results.append([x1, y1, x2, y2])79 80    print("Number of frames cropped: {}".format(len(results)))81    print("First coords: {}".format(results[0]))82    boxes = np.array(results)83    boxes = get_smoothened_boxes(boxes, T=25)84    # results = [[image[y1:y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]85 86    del detector87    return boxes88 89 90def add_black(imgs):91    for i in range(len(imgs)):92        imgs[i] = cv2.vconcat([np.zeros((100, imgs[i].shape[1], 3), dtype=np.uint8), imgs[i], np.zeros((20, imgs[i].shape[1], 3), dtype=np.uint8)])93 94    return imgs95 96 97def preprocess(video_dir="./assets/videos", save_dir="./assets/coords"):98    all_videos = natsorted(glob.glob(os.path.join(video_dir, "*.mp4")))99    for video_path in all_videos:100        video_stream = cv2.VideoCapture(video_path)101 102        # print('Reading video frames...')103        full_frames = []104        while 1:105            still_reading, frame = video_stream.read()106            if not still_reading:107                video_stream.release()108                break109            full_frames.append(frame)110        print("Number of frames available for inference: " + str(len(full_frames)))111        full_frames = add_black(full_frames)112        # print('Face detection running...')113        coords = face_detect(full_frames, pads=(0, 0, 0, 0))114        np.savez_compressed(os.path.join(save_dir, os.path.basename(video_path).split(".")[0]), coords=coords)115 116 117def load_from_npz(video_name, save_dir="./assets/coords"):118    npz = np.load(os.path.join(save_dir, video_name + ".npz"))119    return npz["coords"]120 121 122if __name__ == "__main__":123    preprocess()124