Spaces:
Sleeping
Sleeping
File size: 2,326 Bytes
1dc2504 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | from __future__ import annotations
import argparse
from pathlib import Path
from typing import Dict, List
import pandas as pd
VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv"}
def discover_dataset(root: Path, dataset_name: str) -> List[Dict[str, str]]:
rows: List[Dict[str, str]] = []
for path in root.rglob("*"):
if path.suffix.lower() not in VIDEO_EXTS:
continue
label = 1 if "fake" in str(path).lower() else 0
identity = path.parent.name
rows.append(
{
"dataset": dataset_name,
"video_path": str(path.resolve()),
"label": label,
"identity": identity,
}
)
return rows
def identity_disjoint_split(df: pd.DataFrame, train: float, val: float) -> pd.DataFrame:
identities = sorted(df["identity"].unique())
n = len(identities)
n_train = int(n * train)
n_val = int(n * val)
train_ids = set(identities[:n_train])
val_ids = set(identities[n_train : n_train + n_val])
def split_for_identity(identity: str) -> str:
if identity in train_ids:
return "train"
if identity in val_ids:
return "val"
return "test"
df["split"] = df["identity"].map(split_for_identity)
return df
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--data-root", type=Path, required=True)
parser.add_argument("--out", type=Path, required=True)
parser.add_argument("--train-ratio", type=float, default=0.7)
parser.add_argument("--val-ratio", type=float, default=0.15)
args = parser.parse_args()
rows: List[Dict[str, str]] = []
ffpp = args.data_root / "FaceForensics++"
celeb = args.data_root / "Celeb-DF"
if ffpp.exists():
rows.extend(discover_dataset(ffpp, "FaceForensics++"))
if celeb.exists():
rows.extend(discover_dataset(celeb, "Celeb-DF"))
if not rows:
raise FileNotFoundError("No supported videos found under data root.")
df = pd.DataFrame(rows)
df = identity_disjoint_split(df, args.train_ratio, args.val_ratio)
args.out.parent.mkdir(parents=True, exist_ok=True)
df.to_csv(args.out, index=False)
print(f"Saved metadata: {args.out} ({len(df)} videos)")
if __name__ == "__main__":
main()
|