File size: 5,486 Bytes
4b0b144 | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | #!/usr/bin/env python3
"""
Build a balanced 10k-real / 10k-fake face-image folder with maximum
deepfake-tech variety.
Usage:
python build_face_dataset.py --out_dir ./faces20k --seed 42
"""
import argparse
import os
import pathlib
import random
import shutil
import subprocess
import sys
import zipfile
from collections import defaultdict
import pandas as pd
from tqdm import tqdm
# ----------------------------------------------------------------------
# 1. Edit here to add / remove sources
# ----------------------------------------------------------------------
DATASETS = [
{
"name": "140k",
"slug": "xhlulu/140k-real-and-fake-faces",
"subdirs": {"real": "real", "fake": "fake"},
"fake_label": "stylegan2",
},
{
"name": "deepfake_real",
"slug": "manjilkarki/deepfake-and-real-images",
"subdirs": {"real": "real", "fake": "fake"},
"fake_label": "pggan_stylegan_mix",
},
{
"name": "dfdc_f150",
"slug": "sciarrilli/dfdc-f150",
"subdirs": {"real": "real", "fake": "fake"},
"fake_label": "dfdc_swaps",
},
{
"name": "faceforensics_imgs",
"slug": "greatgamedota/faceforensics",
"subdirs": {"real": "real", "fake": "fake"},
"fake_label": "ffpp_swaps",
},
]
TARGET_PER_CLASS = 10_000
# ----------------------------------------------------------------------
def kaggle_download(slug: str, dest: pathlib.Path) -> pathlib.Path:
"""Download <slug> to dest/. Returns path of the zip."""
dest.mkdir(parents=True, exist_ok=True)
zip_path = dest / f"{slug.split('/')[-1]}.zip"
if zip_path.exists():
return zip_path
print(f"Downloading {slug} …")
subprocess.run(
["kaggle", "datasets", "download", "-d", slug, "-p", str(dest), "--quiet"],
check=True,
)
return zip_path
def extract(zip_path: pathlib.Path, dest: pathlib.Path) -> pathlib.Path:
"""Unzip if needed. Returns extraction dir."""
extract_dir = dest / zip_path.stem
if extract_dir.exists():
return extract_dir
print(f"Extracting {zip_path.name} …")
with zipfile.ZipFile(zip_path) as zf:
zf.extractall(path=extract_dir)
return extract_dir
def glob_images(root: pathlib.Path, pattern: str):
return list(root.glob(pattern)) + list(root.glob(pattern.replace("jpg", "png")))
def main(out_dir: pathlib.Path, seed: int):
random.seed(seed)
temp_root = out_dir / "_raw"
real_pool, fake_pool = [], []
fake_source_tag = {} # path -> dataset tag
# ------------------------------------------------------------------
# 2. Pull sources
# ------------------------------------------------------------------
for ds in DATASETS:
zip_path = kaggle_download(ds["slug"], temp_root)
extract_dir = extract(zip_path, temp_root)
real_dir = extract_dir / ds["subdirs"]["real"]
fake_dir = extract_dir / ds["subdirs"]["fake"]
real_pool += glob_images(real_dir, "**/*.jpg")
fakes = glob_images(fake_dir, "**/*.jpg")
fake_pool += fakes
for fp in fakes:
fake_source_tag[str(fp)] = ds["fake_label"]
# sanity check
if len(real_pool) < TARGET_PER_CLASS or len(fake_pool) < TARGET_PER_CLASS:
print("Not enough images – add another dataset.", file=sys.stderr)
sys.exit(1)
# ------------------------------------------------------------------
# 3. Sample
# ------------------------------------------------------------------
random.shuffle(real_pool)
random.shuffle(fake_pool)
# try to spread fake quota equally over sources
per_source_quota = TARGET_PER_CLASS // len(DATASETS)
selected_fake = []
taken = defaultdict(int)
for fp in fake_pool:
tag = fake_source_tag[str(fp)]
if taken[tag] < per_source_quota:
selected_fake.append(fp)
taken[tag] += 1
if len(selected_fake) == TARGET_PER_CLASS:
break
# top-up if we’re short (some sets too small)
if len(selected_fake) < TARGET_PER_CLASS:
needed = TARGET_PER_CLASS - len(selected_fake)
selected_fake += fake_pool[len(selected_fake) : len(selected_fake) + needed]
selected_real = real_pool[:TARGET_PER_CLASS]
# ------------------------------------------------------------------
# 4. Copy to final tree + manifest
# ------------------------------------------------------------------
for cls in ("real", "fake"):
(out_dir / cls).mkdir(parents=True, exist_ok=True)
manifest_rows = []
def copy_files(file_list, cls):
for src in tqdm(file_list, desc=f"Copying {cls}"):
dst = out_dir / cls / src.name
shutil.copy(src, dst)
manifest_rows.append(
{
"filename": dst.name,
"label": cls,
"source": fake_source_tag.get(str(src), "n/a"),
}
)
copy_files(selected_real, "real")
copy_files(selected_fake, "fake")
pd.DataFrame(manifest_rows).to_csv(out_dir / "manifest.csv", index=False)
print("Done →", out_dir)
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument(
"--out_dir", default="faces20k", type=pathlib.Path, help="destination folder"
)
p.add_argument("--seed", default=42, type=int)
args = p.parse_args()
main(args.out_dir, args.seed)
|