prepare_hf_dataset.py

"""
Prepare HuggingFace dataset for ManipShield HF training branch.

Downloads:
  - ash12321/nano-banana-pro-generated-1k  ->fake images (AI-generated)
  - huggan/wikiart                          ->real images (artwork, sampled)

Splits 80/20 into:
  data_hf/
    train/real/   train/fake/
    val/real/     val/fake/

Usage:
  python prepare_hf_dataset.py
  python prepare_hf_dataset.py --fake-limit 1000 --real-limit 1000 --val-split 0.2
"""

import argparse
import os
import shutil
import random
from pathlib import Path

from PIL import Image


ROOT = Path(__file__).parent
DATA_HF = ROOT / "data_hf"


def save_image(img, dest: Path, idx: int) -> bool:
    """Save a PIL Image or HF image object to dest, return True on success."""
    dest.parent.mkdir(parents=True, exist_ok=True)
    try:
        if not isinstance(img, Image.Image):
            img = img.convert("RGB") if hasattr(img, "convert") else Image.fromarray(img)
        img = img.convert("RGB")
        img.save(dest)
        return True
    except Exception as e:
        print(f"  [skip] {dest.name}: {e}")
        return False


def download_fake(limit: int) -> Path:
    """Stream nano-banana-pro-generated-1k ->data_hf/all_fake/"""
    from datasets import load_dataset

    out = DATA_HF / "all_fake"
    out.mkdir(parents=True, exist_ok=True)

    existing = len(list(out.glob("*.jpg")))
    if existing >= limit:
        print(f"  [skip] {existing} fake images already in {out}")
        return out

    print(f"Downloading ash12321/nano-banana-pro-generated-1k (limit={limit}) ...")
    ds = load_dataset(
        "ash12321/nano-banana-pro-generated-1k",
        split="train",
        streaming=True,
        trust_remote_code=True,
    )

    count = existing
    for item in ds:
        if count >= limit:
            break
        img = item.get("image") or item.get("img") or list(item.values())[0]
        dest = out / f"fake_{count:04d}.jpg"
        if dest.exists():
            count += 1
            continue
        if save_image(img, dest, count):
            count += 1
            if count % 100 == 0:
                print(f"  fake: {count}/{limit}")

    print(f"  Downloaded {count} fake images ->{out}")
    return out


def download_real(limit: int) -> Path:
    """Stream huggan/wikiart ->data_hf/all_real/ (sampled)"""
    from datasets import load_dataset

    out = DATA_HF / "all_real"
    out.mkdir(parents=True, exist_ok=True)

    existing = len(list(out.glob("*.jpg")))
    if existing >= limit:
        print(f"  [skip] {existing} real images already in {out}")
        return out

    print(f"Downloading huggan/wikiart (sampling {limit} images) ...")
    # wikiart has multiple splits / configs — use default
    try:
        ds = load_dataset(
            "huggan/wikiart",
            split="train",
            streaming=True,
            trust_remote_code=True,
        )
    except Exception:
        # Some HF datasets need the config name
        ds = load_dataset(
            "huggan/wikiart",
            name="default",
            split="train",
            streaming=True,
            trust_remote_code=True,
        )

    count = existing
    # Skip forward by 'existing' to avoid re-downloading
    ds_iter = iter(ds)
    for _ in range(existing):
        try:
            next(ds_iter)
        except StopIteration:
            break

    for item in ds_iter:
        if count >= limit:
            break
        img = item.get("image") or item.get("img") or list(item.values())[0]
        dest = out / f"real_{count:04d}.jpg"
        if save_image(img, dest, count):
            count += 1
            if count % 100 == 0:
                print(f"  real: {count}/{limit}")

    print(f"  Downloaded {count} real images ->{out}")
    return out


def split_and_organize(fake_dir: Path, real_dir: Path, val_split: float):
    """80/20 split ->data_hf/train/{real,fake} and data_hf/val/{real,fake}.
    Auto-balances classes to 1:1 by capping the larger class."""
    for split in ["train", "val"]:
        for cls in ["real", "fake"]:
            (DATA_HF / split / cls).mkdir(parents=True, exist_ok=True)

    fake_imgs = sorted(fake_dir.glob("*.jpg"))
    real_imgs = sorted(real_dir.glob("*.jpg"))

    # Balance to 1:1 ratio
    n = min(len(fake_imgs), len(real_imgs))
    if len(fake_imgs) != len(real_imgs):
        print(f"  [balance] {len(real_imgs)} real / {len(fake_imgs)} fake ->capping to {n}:{n}")
    random.shuffle(fake_imgs)
    random.shuffle(real_imgs)
    fake_imgs = fake_imgs[:n]
    real_imgs = real_imgs[:n]

    for cls, imgs in [("fake", fake_imgs), ("real", real_imgs)]:
        n_val = max(1, int(len(imgs) * val_split))
        val_imgs   = imgs[:n_val]
        train_imgs = imgs[n_val:]

        for img_path in train_imgs:
            dest = DATA_HF / "train" / cls / img_path.name
            if not dest.exists():
                shutil.copy2(img_path, dest)

        for img_path in val_imgs:
            dest = DATA_HF / "val" / cls / img_path.name
            if not dest.exists():
                shutil.copy2(img_path, dest)

        print(f"  {cls}: {len(train_imgs)} train, {len(val_imgs)} val")


def main():
    parser = argparse.ArgumentParser(description="Prepare HF dataset for ManipShield-HF training")
    parser.add_argument("--fake-limit", type=int, default=1000,
                        help="Max fake images to download (default 1000)")
    parser.add_argument("--real-limit", type=int, default=1000,
                        help="Max real images to sample from wikiart (default 1000)")
    parser.add_argument("--val-split", type=float, default=0.2,
                        help="Fraction for validation set (default 0.2)")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    random.seed(args.seed)

    print("=" * 60)
    print("ManipShield HF Dataset Preparation")
    print("=" * 60)

    fake_dir = download_fake(args.fake_limit)
    real_dir = download_real(args.real_limit)

    print("\nSplitting into train/val ...")
    split_and_organize(fake_dir, real_dir, args.val_split)

    # Stats
    for split in ["train", "val"]:
        for cls in ["real", "fake"]:
            n = len(list((DATA_HF / split / cls).glob("*.jpg")))
            print(f"  data_hf/{split}/{cls}: {n} images")

    print("\nDone! Train ManipShield-HF with:")
    print(
        "  cd v1/inceptrix/ml/detection && python train.py \\\n"
        f"    --data_dir {DATA_HF.as_posix()} \\\n"
        "    --epochs 50 --batch_size 4 --num_workers 0 \\\n"
        "    --lr 1e-4 --lora_rank 8 \\\n"
        f"    --checkpoint_dir {(ROOT / 'checkpoints_hf').as_posix()} \\\n"
        "    --patience 10"
    )


if __name__ == "__main__":
    main()