benchmark/datasets_util.py

"""Shared helpers for the benchmark scripts: download a dataset with its built-in loader and find where to evaluate it.

``ensure_dataset`` composes the dataset's Hydra config (``conf/dataset/<modality>/<name>.yaml``), and, if the data is not
on disk yet, constructs its dataset module, whose loader (``prefetch: true``) downloads the data and restructures it into
``<generator>/{0_real,1_fake}``. ``find_eval_root`` then locates the folder to evaluate and picks the ``bench_eval.py``
profile (forensynths / ufd / pergen / pool) from the layout.
"""
import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

IMAGE_EXT = {".jpg", ".jpeg", ".jfif", ".png", ".webp", ".bmp", ".tif", ".tiff"}
VIDEO_EXT = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".gif"}
REAL_NAMES = {"0_real", "real", "reals"}
FAKE_NAMES = {"1_fake", "fake", "fakes", "ai"}
# datasets whose generator grouping / real pairing is special-cased by a bench_eval profile
PROFILE_OVERRIDE = {"cnndetection": "forensynths", "cospy_subset": "forensynths", "spai_ldm": "forensynths",
                    "universialfakedetection": "ufd"}


def dataset_config(name: str, modality: str = "image", root: str | None = None):
    """The composed Hydra config for ``conf/dataset/<modality>/<name>.yaml`` (``root`` overrides ``root_dir``)."""
    from config_utils import compose_config
    conf = ROOT / "conf" / "dataset"
    if not (conf / modality / f"{name}.yaml").is_file():
        other = [m for m in ("image", "video") if (conf / m / f"{name}.yaml").is_file()]
        if other:
            raise SystemExit(f"{name} is a {other[0]} dataset; use the {other[0]} evaluator "
                             f"({'benchmark/bench_eval_video.py' if other[0] == 'video' else 'benchmark/bench_eval.py'})")
        known = sorted(p.stem for p in (conf / modality).glob("*.yaml"))
        raise SystemExit(f"unknown {modality} dataset '{name}'. Available: {', '.join(known)}")
    overrides = [f"dataset@DatasetConfig={modality}/{name}"]
    if root:
        overrides.append(f"DatasetConfig.root_dir={root}")
    return compose_config(overrides)


def has_images(root: Path, exts=IMAGE_EXT) -> bool:
    """True as soon as one file with an extension in ``exts`` (images by default) is found under ``root`` (cheap: stops at
    the first hit)."""
    if not root.is_dir():
        return False
    for _, _, files in os.walk(root):
        if any(os.path.splitext(f)[1].lower() in exts for f in files):
            return True
    return False


def download_dataset(name: str, modality: str = "image", root: str | None = None) -> Path:
    """Run the dataset's loader (download + restructure). Returns the root directory."""
    from dataset import AIGIDatasetModule
    cfg = dataset_config(name, modality, root)
    AIGIDatasetModule(cfg)        # constructing the datasets runs the loader (prefetch: true downloads + restructures)
    return Path(cfg.DatasetConfig.root_dir)


def ensure_dataset(name: str, modality: str = "image", root: str | None = None, download: bool = True) -> Path:
    """The dataset's root directory, downloading it first when no images (or videos) are there yet."""
    cfg = dataset_config(name, modality, root)
    path = Path(cfg.DatasetConfig.root_dir).expanduser()
    if has_images(path, VIDEO_EXT if modality == "video" else IMAGE_EXT):
        return path
    if not download:
        raise SystemExit(f"{name}: no data under {path} (drop --no-download to fetch it, or pass --root)")
    print(f"[dataset] {name}: nothing under {path}, downloading with the built-in loader ...", flush=True)
    return download_dataset(name, modality, root)


def _subdirs(d: Path) -> list[Path]:
    try:
        return sorted(c for c in d.iterdir() if c.is_dir() and not c.name.startswith((".", "_")))
    except OSError:
        return []


def _has_real_fake(d: Path, depth: int = 3) -> bool:
    names = {c.name.lower() for c in d.iterdir() if c.is_dir()} if d.is_dir() else set()
    if names & REAL_NAMES and names & FAKE_NAMES:
        return True
    return depth > 0 and any(_has_real_fake(c, depth - 1) for c in _subdirs(d)[:60])


def find_eval_root(name: str, root: Path) -> tuple[Path, str]:
    """(folder to evaluate, bench_eval profile). Looks in ``root/test``, ``root/test_full`` and ``root`` itself."""
    if name in PROFILE_OVERRIDE:       # fixed layouts handled by their own bench_eval profile (e.g. UFD: fakes and reals in separate folders)
        for cand in (root / "test", root):
            if cand.is_dir() and _subdirs(cand):
                return cand, PROFILE_OVERRIDE[name]
    for cand in (root / "test", root / "test_full", root):
        if not cand.is_dir():
            continue
        names = {c.name.lower() for c in cand.iterdir() if c.is_dir()}
        if names & REAL_NAMES and names & FAKE_NAMES:            # one pooled real/fake set
            return cand, PROFILE_OVERRIDE.get(name, "pool")
        if any(_has_real_fake(g) for g in _subdirs(cand)):       # one folder per generator
            return cand, PROFILE_OVERRIDE.get(name, "pergen")
    raise SystemExit(f"{name}: could not find real/fake folders under {root}; expected "
                     "<root>/[test/]<generator>/{0_real,1_fake} (see docs/BENCHMARKS.md)")