slotread/data/__init__.py

"""Datasets. Parsed examples are cached on disk (JPEG + index.json) under SLOTREAD_CACHE, because
streaming and filtering the sources is slow (AGUVIS: up to hours) and every arm trains on the same
examples."""

import json, os, resource

from PIL import Image

CACHE = os.path.expanduser(os.environ.get("SLOTREAD_CACHE", "~/.cache/slotread"))

# cached examples are opened lazily (PIL holds one descriptor each); tens of thousands of them need a
# higher descriptor limit than the default 1024
_soft, _hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (max(_soft, min(65536, _hard)), _hard))


def load_cache(name):
    path = os.path.join(CACHE, name)
    if not os.path.exists(f"{path}/index.json"):
        return None
    out = []
    for r in json.load(open(f"{path}/index.json")):
        r = dict(r)
        r["image"] = Image.open(f"{path}/{r.pop('file')}")
        out.append(r)
    print(f"  cache: {len(out)} examples from {path}", flush=True)
    return out


def save_cache(name, rows):
    path = os.path.join(CACHE, name)
    os.makedirs(path, exist_ok=True)
    idx = []
    for i, r in enumerate(rows):
        f = f"{i:05d}.jpg"
        r["image"].convert("RGB").save(f"{path}/{f}", "JPEG", quality=88)
        idx.append({k: v for k, v in r.items() if k != "image"} | {"file": f})
    json.dump(idx, open(f"{path}/index.json", "w"))
    print(f"  cached {len(rows)} examples to {path}", flush=True)