import os from concurrent.futures import ThreadPoolExecutor from functools import partial from typing import Callable import numpy as np import pandas as pd from PIL import Image, ImageFile from tqdm import tqdm from src.utils import logger from src.utils.logger import print from .base import BaseDataset ImageFile.LOAD_TRUNCATED_IMAGES = True class DeepfakeDataset(BaseDataset): """ DeepfakeDataset is any dataset that follows this structure: ... / / / / - Name of the dataset, e.g. FF, CDF, DFD, DFDC... - Name of the source, e.g. real, fake, or any name of generator, e.g. FaceSwap, Face2Face... - Name of the video, e.g. 000, 000_003, ... - Any name of the frame, e.g. 000001.jpg, 000002.jpg, ... Labels are automatically created from such that: - If has "real" substring, the label is 0 - Otherwise, the label is 1 """ def __init__( self, files_with_paths: list[str] | dict[str, list[str]], preprocess: None | Callable = None, augmentations: None | Callable = None, shuffle: bool = False, # Shuffles the dataset once binary: bool = False, limit_files: None | int = None, load_pairs: bool = False, ): files = [] labels = [] logger.print_info("Loading files") if binary: label2name = {0: "real", 1: "fake"} else: raise NotImplementedError("Only binary classification is supported now") source2label = {v: k for k, v in label2name.items()} self.label2name = label2name dataset2files = None if isinstance(files_with_paths, dict): dataset2files_with_paths = files_with_paths.copy() dataset2files = {dataset_name: [] for dataset_name in dataset2files_with_paths.keys()} files_with_paths = [item for sublist in files_with_paths.values() for item in sublist] max_workers = min(64, os.cpu_count()) for file_with_paths in sorted(set(files_with_paths)): with open(file_with_paths, "r") as f: paths = f.readlines() paths = [path.strip() for path in paths] # If files do not exist, append root of 'file' to each path root = file_with_paths.rsplit("/", 1)[0] def process_path(root, path): if not os.path.exists(path): path = f"{root}/{path}" assert os.path.exists(path), f"File not found: {path}" return path with ThreadPoolExecutor(max_workers) as executor: process_with_root = partial(process_path, root) paths = list( tqdm( executor.map(process_with_root, paths), total=len(paths), desc=f"Processing paths in {file_with_paths}", leave=True, ) ) files.extend(paths) if dataset2files is not None: for dataset_name, files_with_paths in dataset2files_with_paths.items(): if file_with_paths in files_with_paths: dataset2files[dataset_name].extend(paths) # Remove duplicate paths files = np.unique(files).tolist() # Limit the number of files if limit_files is not None: files = self.limit_files(files, limit_files) # Get labels from paths for path in files: source = self.get_source_from_file(path) if binary: if "real" in source: source = "real" else: source = "fake" label = source2label[source] labels.append(label) logger.print_info("Files loaded") super().__init__(files, labels, preprocess, augmentations, shuffle, dataset2files) self.source2uid = self._source2uid() self.video_path2uid = self._video_path2uid() self.file2index = {f: i for i, f in enumerate(self.files)} def limit_files(self, files: list[str], limit: int) -> list[str]: """Limits number of files by considering unique videos""" # Select unique videos video_paths = [self.get_video_path(file) for file in files] unique_videos = list(np.unique(video_paths)) # For each video, select files video2files = {video: [] for video in unique_videos} for file, video in zip(files, video_paths): video2files[video].append(file) # Shuffle videos with fixed seed np.random.RandomState(42).shuffle(unique_videos) # Select files from shuffled videos selected_files = [] for video in unique_videos: selected_files.extend(video2files[video]) if len(selected_files) >= limit: break return selected_files[:limit] def _source2uid(self) -> dict[str, int]: sources = [self.get_source_from_file(file) for file in self.files] sources = np.unique(sources) assert any("real" in g for g in sources), "No real source found" sources = [str(g) for g in sources] # Map all real sources to 0 and fake sources to 1, 2, 3, ... real_sources = [g for g in sources if "real" in g] fake_sources = [g for g in sources if "real" not in g] source2uid = dict.fromkeys(real_sources, 0) for i, s in enumerate(fake_sources, start=1): source2uid[s] = i return source2uid def _video_path2uid(self) -> dict[str, int]: video_paths = [self.get_video_path(file) for file in self.files] unique_videos = list(np.unique(video_paths)) return {video: i for i, video in enumerate(unique_videos)} @staticmethod def get_frame_from_file(file_path: str) -> str: # ... / / / / # returns return file_path.split("/")[-1] @staticmethod def get_video_from_file(file_path: str) -> str: # ... / / / / # returns return file_path.split("/")[-2] @staticmethod def get_source_from_file(file_path: str) -> str: # ... / / / / # returns return file_path.split("/")[-3] @staticmethod def get_dataset_from_file(file_path: str) -> str: # ... / / / / # returns return file_path.split("/")[-4] @staticmethod def get_video_path(file_path: str) -> str: # ... / / / / # file_path[::-1].find("/") finds the last occurrence of "/" # returns ...//// return file_path[: -file_path[::-1].find("/")] def get_class_names(self) -> dict[int, str]: return self.label2name def print_statistics(self): super().print_statistics() video_paths = [self.get_video_path(file) for file in self.files] files_by_dataset = [self.get_dataset_from_file(file) for file in self.files] print(f"Total number of frames: {len(self.files)}") print(f"Total number of videos: {len(set(video_paths))}") # For each dataset, print number of frames and videos df = pd.DataFrame({"dataset": files_by_dataset, "video": video_paths}) for dataset in df["dataset"].unique(): dataset_df = df[df["dataset"] == dataset] videos_count = dataset_df["video"].nunique() frames_count = len(dataset_df) print(f"Dataset: {dataset}, videos: {videos_count}, frames: {frames_count}") def __getitem__(self, idx): path = self.files[idx] image = Image.open(path) source = self.get_source_from_file(path) video_path = self.get_video_path(path) label = self.labels[idx] # Apply augmentations defined in from config.Augmentations if self.augmentations is not None: image = self.augmentations(image) # Apply preprocessing defined by the model input requirements if self.preprocess is not None: image = self.preprocess(image) output = { "idx": idx, "image": image, "label": label, "path": path, "video": self.get_video_from_file(path), "source_uid": self.source2uid[source], "frame": self.get_frame_from_file(path), "video_uid": self.video_path2uid[video_path], } return output