Spaces:
Runtime error
Runtime error
| """ | |
| Video Face Manipulation Detection Through Ensemble of CNNs | |
| Image and Sound Processing Lab - Politecnico di Milano | |
| Nicolò Bonettini | |
| Edoardo Daniele Cannas | |
| Sara Mandelli | |
| Luca Bondi | |
| Paolo Bestagini | |
| """ | |
| from typing import List | |
| import albumentations as A | |
| import pandas as pd | |
| from albumentations.pytorch import ToTensorV2 | |
| from .data import FrameFaceIterableDataset, get_iterative_real_fake_idxs | |
| class FrameFaceTripletIterableDataset(FrameFaceIterableDataset): | |
| def __init__(self, | |
| roots: List[str], | |
| dfs: List[pd.DataFrame], | |
| size: int, | |
| scale: str, | |
| num_triplets: int = -1, | |
| transformer: A.BasicTransform = ToTensorV2(), | |
| seed: int = None): | |
| """ | |
| :param roots: List of root folders for frames cache | |
| :param dfs: List of DataFrames of cached frames with 'bb' column as array of 4 elements (left,top,right,bottom) | |
| and 'label' column | |
| :param size: face size | |
| :param num_triplets: number of samples for the dataset | |
| :param idxs: sampling indexes triplets (each element is a key for anchor, positive, negative) | |
| :param scale: Rescale the face to the given size, preserving the aspect ratio. | |
| If false crop around center to the given size | |
| :param transformer: | |
| :param seed: | |
| """ | |
| super(FrameFaceTripletIterableDataset, self).__init__( | |
| roots=roots, | |
| dfs=dfs, | |
| size=size, | |
| scale=scale, | |
| num_samples=num_triplets * 3, | |
| transformer=transformer, | |
| seed=seed | |
| ) | |
| self.num_triplet_couples = self.num_samples // 6 | |
| self.num_triplets = self.num_triplet_couples * 2 | |
| self.num_samples = self.num_triplets * 3 | |
| def __len__(self): | |
| return self.num_triplets | |
| def __iter__(self): | |
| random_fake_idxs, random_real_idxs = get_iterative_real_fake_idxs( | |
| df_real=self.df_real, | |
| df_fake=self.df_fake, | |
| num_samples=self.num_samples, | |
| seed0=self.seed0 | |
| ) | |
| while len(random_fake_idxs) >= 3 and len(random_real_idxs) >= 3: | |
| a = self._get_face(random_fake_idxs.pop())[0] | |
| p = self._get_face(random_fake_idxs.pop())[0] | |
| n = self._get_face(random_real_idxs.pop())[0] | |
| yield a, p, n | |
| a = self._get_face(random_real_idxs.pop())[0] | |
| p = self._get_face(random_real_idxs.pop())[0] | |
| n = self._get_face(random_fake_idxs.pop())[0] | |
| yield a, p, n | |