| from __future__ import annotations |
|
|
| import argparse |
| import io |
| import json |
| import os |
| import random |
| import zipfile |
| from concurrent.futures import ThreadPoolExecutor |
|
|
| import numpy as np |
| import requests |
| from PIL import Image |
|
|
| ANN_URL = "http://images.cocodataset.org/annotations/annotations_trainval2014.zip" |
|
|
| def csr(img, size): |
| img = img.convert("RGB") |
| w, h = img.size |
| s = min(w, h) |
| l, t = (w - s) // 2, (h - s) // 2 |
| return np.asarray(img.crop((l, t, l + s, t + s)).resize((size, size), Image.BICUBIC), dtype=np.uint8) |
|
|
| def fetch_one(item, size): |
| url, cap = item |
| for _ in range(3): |
| try: |
| r = requests.get(url, timeout=15) |
| if r.status_code == 200: |
| return csr(Image.open(io.BytesIO(r.content)), size), cap |
| except Exception: |
| pass |
| return None |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--work", default="/root/v6cache") |
| ap.add_argument("--n", type=int, default=5000) |
| ap.add_argument("--size", type=int, default=256) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--workers", type=int, default=48) |
| args = ap.parse_args() |
| os.makedirs(args.work, exist_ok=True) |
|
|
| out_path = os.path.join(args.work, "eval_256.npz") |
| if os.path.exists(out_path): |
| print(f"[eval-set] already exists at {out_path}", flush=True) |
| return |
|
|
| ann_path = os.path.join(args.work, "captions_val2014.json") |
| if not os.path.exists(ann_path): |
| print("[eval-set] downloading annotations", flush=True) |
| z = os.path.join(args.work, "ann.zip") |
| with requests.get(ANN_URL, stream=True, timeout=120) as r: |
| with open(z, "wb") as f: |
| for chunk in r.iter_content(1 << 20): |
| f.write(chunk) |
| with zipfile.ZipFile(z) as zf: |
| with zf.open("annotations/captions_val2014.json") as src, open(ann_path, "wb") as dst: |
| dst.write(src.read()) |
| os.remove(z) |
|
|
| ann = json.load(open(ann_path)) |
| url_by_id = {im["id"]: im["coco_url"] for im in ann["images"]} |
| cap_by_id = {} |
| for a in ann["annotations"]: |
| cap_by_id.setdefault(a["image_id"], a["caption"]) |
| items = [(url_by_id[i], cap_by_id[i]) for i in cap_by_id if i in url_by_id] |
| random.Random(args.seed).shuffle(items) |
| print(f"[eval-set] {len(items)} val2014 pairs available, target {args.n}", flush=True) |
|
|
| imgs, caps = [], [] |
| pool = ThreadPoolExecutor(max_workers=args.workers) |
| idx, batch = 0, 64 |
| while len(imgs) < args.n and idx < len(items): |
| chunk = items[idx:idx + batch] |
| idx += batch |
| results = [r for r in pool.map(lambda it: fetch_one(it, args.size), chunk) if r is not None] |
| for a, c in results: |
| imgs.append(a); caps.append(c) |
| if idx % (batch * 20) == 0: |
| print(f"[eval-set] {len(imgs)}/{args.n}", flush=True) |
|
|
| imgs = np.stack(imgs[:args.n]) |
| caps = np.array(caps[:args.n], dtype=object) |
| np.savez(out_path, images=imgs, captions=caps) |
| print(f"[eval-set] DONE {imgs.shape} -> {out_path}", flush=True) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
