svjack/LatentSync

1# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import mediapipe as mp16import cv217from decord import VideoReader18from einops import rearrange19import os20import numpy as np21import torch22import tqdm23from eval.fvd import compute_our_fvd24 25 26class FVD:27    def __init__(self, resolution=(224, 224)):28        self.face_detector = mp.solutions.face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.5)29        self.resolution = resolution30 31    def detect_face(self, image):32        height, width = image.shape[:2]33        # Process the image and detect faces.34        results = self.face_detector.process(image)35 36        if not results.detections:  # Face not detected37            raise Exception("Face not detected")38 39        detection = results.detections[0]  # Only use the first face in the image40        bounding_box = detection.location_data.relative_bounding_box41        xmin = int(bounding_box.xmin * width)42        ymin = int(bounding_box.ymin * height)43        face_width = int(bounding_box.width * width)44        face_height = int(bounding_box.height * height)45 46        # Crop the image to the bounding box.47        xmin = max(0, xmin)48        ymin = max(0, ymin)49        xmax = min(width, xmin + face_width)50        ymax = min(height, ymin + face_height)51        image = image[ymin:ymax, xmin:xmax]52 53        return image54 55    def detect_video(self, video_path, real: bool = True):56        vr = VideoReader(video_path)57        video_frames = vr[20:36].asnumpy()  # Use one frame per second58        vr.seek(0)  # avoid memory leak59        faces = []60        for frame in video_frames:61            face = self.detect_face(frame)62            face = cv2.resize(face, (self.resolution[1], self.resolution[0]), interpolation=cv2.INTER_AREA)63            faces.append(face)64 65        if len(faces) != 16:66            return None67        faces = np.stack(faces, axis=0)  # (f, h, w, c)68        faces = torch.from_numpy(faces)69        return faces70 71 72def eval_fvd(real_videos_dir, fake_videos_dir):73    fvd = FVD()74    real_features_list = []75    fake_features_list = []76    for file in tqdm.tqdm(os.listdir(fake_videos_dir)):77        if file.endswith(".mp4"):78            real_video_path = os.path.join(real_videos_dir, file.replace("_out.mp4", ".mp4"))79            fake_video_path = os.path.join(fake_videos_dir, file)80            real_features = fvd.detect_video(real_video_path, real=True)81            fake_features = fvd.detect_video(fake_video_path, real=False)82            if real_features is None or fake_features is None:83                continue84            real_features_list.append(real_features)85            fake_features_list.append(fake_features)86 87    real_features = torch.stack(real_features_list) / 255.088    fake_features = torch.stack(fake_features_list) / 255.089    print(compute_our_fvd(real_features, fake_features, device="cpu"))90 91 92if __name__ == "__main__":93    real_videos_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/segmented/cross"94    fake_videos_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/segmented/latentsync_cross"95 96    eval_fvd(real_videos_dir, fake_videos_dir)97