training/prepare_defactify_sample.py

from __future__ import annotations

import argparse
import hashlib
import shutil
import sys
from pathlib import Path
from typing import Dict

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from training.media_manifest import MediaRecord, sha256_file, stable_group, write_jsonl


LABEL_TO_FOLDER = {
    0: "real_camera",
    1: "ai_generated",
}

GENERATOR_NAMES = {
    0: "authentic",
    1: "stable-diffusion-2.1",
    2: "stable-diffusion-xl",
    3: "stable-diffusion-3",
    4: "dall-e-3",
    5: "midjourney-6",
}

V4_GENERATOR_SPLITS = {
    1: "train",
    2: "train",
    3: "tuning",
    4: "calibration",
    5: "locked_test",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Download a balanced local sample from the Defactify real-vs-AI image dataset."
    )
    parser.add_argument(
        "--dataset-name",
        default="Rajarshi-Roy-research/Defactify_Image_Dataset",
        help="Hugging Face dataset id.",
    )
    parser.add_argument(
        "--output-dir",
        default="training/data/defactify_sample",
        help="Where to write ImageFolder-style files.",
    )
    parser.add_argument(
        "--max-per-label",
        type=int,
        default=800,
        help="Maximum images per label per split. Start small on a laptop.",
    )
    parser.add_argument(
        "--splits",
        nargs="+",
        default=["train", "validation", "test"],
        help="Dataset splits to export.",
    )
    parser.add_argument(
        "--clean-output",
        action="store_true",
        help="Remove the selected output split folders before exporting. Use this when changing sample sizes.",
    )
    parser.add_argument(
        "--split-policy",
        choices=("source", "generator-heldout-v4"),
        default="generator-heldout-v4",
        help="Use generator-separated v4 splits or retain the dataset's source split.",
    )
    parser.add_argument(
        "--dataset-license",
        default="license-review-required",
        help="License identifier recorded in the manifest. Confirm it before public redistribution.",
    )
    return parser.parse_args()


def main() -> None:
    from datasets import load_dataset

    args = parse_args()
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    target_counts: Dict[tuple[int, str], int] = {}

    if args.clean_output:
        target_names = set(args.splits)
        if args.split_policy == "generator-heldout-v4":
            target_names.update({"train", "tuning", "calibration", "locked_test"})
        for target_name in target_names:
            target_dir = output_dir / target_name
            if target_dir.exists():
                shutil.rmtree(target_dir)

    manifest_records: list[MediaRecord] = []
    for split in args.splits:
        print(f"Preparing split: {split}")
        split_dir = output_dir / split
        dataset = load_dataset(args.dataset_name, split=split, streaming=True)
        counts: Dict[int, int] = {label: 0 for label in LABEL_TO_FOLDER}

        for example in dataset:
            raw_label = example.get("Label_A")
            if raw_label not in LABEL_TO_FOLDER:
                continue
            raw_generator = _integer_label(example.get("Label_B"), default=0 if raw_label == 0 else -1)
            image = example.get("Image")
            if image is None:
                continue
            caption = _first_text(example, "Caption", "caption", "Prompt", "prompt", "Text", "text")
            source_id = _first_text(example, "id", "ID", "source_id", "filename", "File_Name")
            pixel_group = hashlib.sha256(image.convert("RGB").tobytes()).hexdigest()
            source_group_key = caption or source_id or pixel_group
            target_split = _target_split(
                source_split=split,
                raw_label=raw_label,
                raw_generator=raw_generator,
                source_group_key=source_group_key,
                split_policy=args.split_policy,
            )
            if target_split is None:
                continue
            count_key = (raw_label, target_split)
            target_count = target_counts.get(count_key, 0)
            if target_count >= args.max_per_label:
                continue

            folder = output_dir / target_split / LABEL_TO_FOLDER[raw_label]
            folder.mkdir(parents=True, exist_ok=True)
            index = target_count
            path = folder / f"{target_split}_{LABEL_TO_FOLDER[raw_label]}_{index:06d}.jpg"
            try:
                image.convert("RGB").save(path, format="JPEG", quality=92)
            except Exception as exc:
                print(f"Skipped one image: {exc}")
                continue

            counts[raw_label] += 1
            target_counts[count_key] = target_count + 1
            group = stable_group(source_group_key)
            manifest_records.append(
                MediaRecord(
                    path=path.relative_to(output_dir).as_posix(),
                    sha256=sha256_file(path),
                    media_type="image",
                    class_label=LABEL_TO_FOLDER[raw_label],
                    source=args.dataset_name,
                    license=args.dataset_license,
                    generator_or_editor=GENERATOR_NAMES.get(raw_generator, f"label-b-{raw_generator}"),
                    parent_media=None,
                    transformation="jpeg_quality_92_export",
                    semantic_category="unspecified",
                    source_group=group,
                    split=target_split,
                )
            )
            if sum(counts.values()) % 100 == 0:
                print(f"  saved {counts}")

        print(f"Finished {split}: {counts}")

    write_jsonl(output_dir / "manifest.v4.jsonl", manifest_records)
    print(f"Wrote {len(manifest_records)} manifest records with Label_B generator identities.")
    print(f"Done. Dataset is at: {output_dir}")


def _integer_label(value: object, default: int) -> int:
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


def _first_text(example: dict, *keys: str) -> str:
    for key in keys:
        value = example.get(key)
        if value is not None and str(value).strip():
            return str(value).strip()
    return ""


def _target_split(
    *,
    source_split: str,
    raw_label: int,
    raw_generator: int,
    source_group_key: str,
    split_policy: str,
) -> str | None:
    if split_policy == "source":
        return source_split
    group_split = _group_split(source_group_key)
    if raw_label == 1:
        generator_split = V4_GENERATOR_SPLITS.get(raw_generator)
        # Defactify repeats captions across every generator. Retain a generated
        # variant only when its held-out family belongs to the caption's one
        # assigned split; otherwise the same semantic source would leak across
        # train, tuning, calibration, and locked test.
        return generator_split if generator_split == group_split else None
    return group_split


def _group_split(source_group_key: str) -> str:
    bucket = int(hashlib.sha256(source_group_key.encode("utf-8")).hexdigest()[:8], 16) % 100
    if bucket < 70:
        return "train"
    if bucket < 80:
        return "tuning"
    if bucket < 90:
        return "calibration"
    return "locked_test"


if __name__ == "__main__":
    main()