# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch>=2.1",
#     "datasets>=2.18",
#     "pycocotools",
#     "Pillow",
#     "numpy",
# ]
# ///
import os, io
from PIL import Image
import datasets as hfds

data_dir = "/data"
if not os.path.isdir(data_dir):
    data_dir = "biglam/loc_beyond_words"
ds = hfds.load_dataset(data_dir, split="validation")
examples = list(ds)
print("loaded", len(examples))
ex = examples[0]
print("keys:", list(ex.keys()), "w/h:", ex["width"], ex["height"])
objs = ex["objects"]
print("objects type:", type(objs), "len:", len(objs))
o = objs[0]
print("obj keys:", list(o.keys()))
print("category_id:", repr(o["category_id"]), "type:", type(o["category_id"]).__name__)
print("bbox:", o["bbox"], "type:", type(o["bbox"]).__name__)
img = ex["image"]
print("image type:", type(img).__name__)
im = img.size if not isinstance(img, dict) else Image.open(io.BytesIO(img["bytes"])).size
print("img size:", im)
from collections import Counter
cats = Counter()
for e in examples[:200]:
    for o in e["objects"]:
        cats[repr(o["category_id"])] += 1
print("category_id value distribution sample:", dict(cats))