File size: 4,557 Bytes
e465a2f | 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 | """
fetch_coco_subset.py - build train and eval sets from sayakpaul/coco-30-val-2014.
The stream of 30K (image, caption) rows is split by row index:
rows 0..EVAL_N-1 -> eval set (real images for FID + captions for generation)
rows TRAIN_START..end -> training candidates
Training images whose content hash also appears in the eval set are dropped,
so the model is never trained on an image it is evaluated against (the same
image can appear in multiple rows with different captions). Training rows are
also deduped against each other by image hash.
Outputs (in --out dir):
coco_train.npz images uint8 (N, 64, 64, 3), center-cropped
coco_train.captions.json N caption strings
eval_real/real_00000.jpg... EVAL_N real images, 256x256 center-crop, q95
eval_captions.json EVAL_N caption strings (for generation + CLIP)
Usage:
python fetch_coco_subset.py --out ../pm-work
# much faster: pre-download the 10 parquet shards, then
python fetch_coco_subset.py --out ../pm-work --parquet-dir ../pm-work/shards
"""
import argparse
import glob
import hashlib
import json
import os
import numpy as np
from datasets import load_dataset
from PIL import Image
EVAL_N = 5000
TRAIN_START = 10000 # comfortable gap after the eval rows
TRAIN_RES = 64
EVAL_RES = 256
def center_crop_resize(img: Image.Image, size: int) -> Image.Image:
w, h = img.size
side = min(w, h)
left, top = (w - side) // 2, (h - side) // 2
return img.crop((left, top, left + side, top + side)).resize(
(size, size), Image.BICUBIC
)
def img_hash(img: Image.Image) -> str:
"""Content hash of the image, robust to caption-duplicated rows."""
return hashlib.md5(
img.convert("RGB").resize((64, 64), Image.BILINEAR).tobytes()
).hexdigest()
def main():
p = argparse.ArgumentParser()
p.add_argument("--out", required=True)
p.add_argument("--parquet-dir", default=None,
help="dir with the dataset's parquet shards already downloaded "
"(sorted filename order == hub streaming row order)")
args = p.parse_args()
real_dir = os.path.join(args.out, "eval_real")
os.makedirs(real_dir, exist_ok=True)
if args.parquet_dir:
files = sorted(glob.glob(os.path.join(args.parquet_dir, "*.parquet")))
assert files, f"no parquet files in {args.parquet_dir}"
ds = load_dataset("parquet", data_files=files, split="train", streaming=True)
else:
ds = load_dataset("sayakpaul/coco-30-val-2014", split="train", streaming=True)
eval_hashes = set()
eval_captions = []
train_images, train_captions = [], []
train_hashes = set()
skipped_eval_overlap = skipped_dupe = 0
for i, row in enumerate(ds):
img, caption = row["image"], row["caption"]
if i < EVAL_N:
eval_hashes.add(img_hash(img))
eval_captions.append(caption)
center_crop_resize(img.convert("RGB"), EVAL_RES).save(
os.path.join(real_dir, f"real_{i:05d}.jpg"), quality=95
)
if (i + 1) % 1000 == 0:
print(f" eval: {i + 1}/{EVAL_N}", flush=True)
continue
if i < TRAIN_START:
continue
h = img_hash(img)
if h in eval_hashes:
skipped_eval_overlap += 1
continue
if h in train_hashes:
skipped_dupe += 1
continue
train_hashes.add(h)
thumb = center_crop_resize(img.convert("RGB"), TRAIN_RES)
train_images.append(np.array(thumb, dtype=np.uint8))
train_captions.append(caption)
if len(train_images) % 2000 == 0:
print(f" train: {len(train_images)} (row {i})", flush=True)
images = np.stack(train_images)
np.savez_compressed(os.path.join(args.out, "coco_train.npz"), images=images)
with open(os.path.join(args.out, "coco_train.captions.json"), "w", encoding="utf-8") as f:
json.dump(train_captions, f)
with open(os.path.join(args.out, "eval_captions.json"), "w", encoding="utf-8") as f:
json.dump(eval_captions, f)
print(f"\neval: {len(eval_captions)} pairs @ {EVAL_RES}x{EVAL_RES} -> {real_dir}")
print(f"train: {len(train_captions)} pairs @ {TRAIN_RES}x{TRAIN_RES} "
f"(skipped: {skipped_eval_overlap} eval-overlap, {skipped_dupe} dupes)")
if __name__ == "__main__":
main()
|