103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import mediapipe as mp
|
|
import cv2
|
|
from decord import VideoReader
|
|
from einops import rearrange
|
|
import os
|
|
import numpy as np
|
|
import torch
|
|
import tqdm
|
|
from eval.fvd import compute_our_fvd
|
|
|
|
|
|
class FVD:
|
|
def __init__(self, resolution=(224, 224)):
|
|
self.face_detector = mp.solutions.face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.5)
|
|
self.resolution = resolution
|
|
|
|
def detect_face(self, image):
|
|
height, width = image.shape[:2]
|
|
# Process the image and detect faces.
|
|
results = self.face_detector.process(image)
|
|
|
|
if not results.detections: # Face not detected
|
|
raise Exception("Face not detected")
|
|
|
|
detection = results.detections[0] # Only use the first face in the image
|
|
bounding_box = detection.location_data.relative_bounding_box
|
|
xmin = int(bounding_box.xmin * width)
|
|
ymin = int(bounding_box.ymin * height)
|
|
face_width = int(bounding_box.width * width)
|
|
face_height = int(bounding_box.height * height)
|
|
|
|
# Crop the image to the bounding box.
|
|
xmin = max(0, xmin)
|
|
ymin = max(0, ymin)
|
|
xmax = min(width, xmin + face_width)
|
|
ymax = min(height, ymin + face_height)
|
|
image = image[ymin:ymax, xmin:xmax]
|
|
|
|
return image
|
|
|
|
def detect_video(self, video_path):
|
|
vr = VideoReader(video_path)
|
|
video_frames = vr[20:36].asnumpy() # Use one frame per second
|
|
vr.seek(0) # avoid memory leak
|
|
faces = []
|
|
for frame in video_frames:
|
|
face = self.detect_face(frame)
|
|
face = cv2.resize(face, (self.resolution[1], self.resolution[0]), interpolation=cv2.INTER_AREA)
|
|
faces.append(face)
|
|
|
|
if len(faces) != 16:
|
|
return None
|
|
faces = np.stack(faces, axis=0) # (f, h, w, c)
|
|
faces = torch.from_numpy(faces)
|
|
return faces
|
|
|
|
def detect_videos(self, videos_dir: str):
|
|
videos_list = []
|
|
|
|
if videos_dir.endswith(".mp4"):
|
|
video_faces = self.detect_video(videos_dir)
|
|
if video_faces is None:
|
|
raise RuntimeError("No face detected")
|
|
videos_list.append(video_faces)
|
|
else:
|
|
for file in tqdm.tqdm(os.listdir(videos_dir)):
|
|
if file.endswith(".mp4"):
|
|
video_path = os.path.join(videos_dir, file)
|
|
video_faces = self.detect_video(video_path)
|
|
if video_faces is None:
|
|
raise RuntimeError("No face detected")
|
|
videos_list.append(video_faces)
|
|
|
|
videos_list = torch.stack(videos_list) / 255.0
|
|
return videos_list
|
|
|
|
|
|
def eval_fvd(real_videos_dir: str, fake_videos_dir: str):
|
|
fvd = FVD()
|
|
real_videos = fvd.detect_videos(real_videos_dir)
|
|
fake_videos = fvd.detect_videos(fake_videos_dir)
|
|
print(compute_our_fvd(real_videos, fake_videos, device="cpu"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
real_videos_dir = "dir1"
|
|
fake_videos_dir = "dir2"
|
|
eval_fvd(real_videos_dir, fake_videos_dir)
|