The PerceptionBench dataset now offers a structured way to test vision-language models on specific tasks like OCR, counting, and depth estimation.
Setting up the environment
The tutorial begins by configuring a Google Colab-compatible workspace. It installs necessary libraries for data handling, numerical analysis, and image processing. The script sets up Matplotlib for visualisation and prepares an output directory to store results.
Configuration variables define the dataset repository, the split to use, and the number of samples per capability category. It also specifies the backend for evaluation, which can be a blind prior, an OpenAI-compatible API, or a local Hugging Face model.
import os, sys, io, re, json, time, math, base64, random, hashlib, subprocess, warnings
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
warnings.filterwarnings("ignore")
CFG = dict(
REPO = "moonshotai/PerceptionBench",
SPLIT = "train",
N_PER_CATEGORY = 12,
MAX_SCAN = 1200,
SEED = 0,
LOAD_MODE = "stream",
BACKEND = "blind",
API_BASE = os.environ.get("PB_API_BASE", "https://api.openai.com/v1"),
API_KEY = os.environ.get("PB_API_KEY", ""),
API_MODEL = os.environ.get("PB_API_MODEL", "gpt-4o-mini"),
API_WORKERS = 4,
API_MAX_TOKENS = 512,
LOCAL_MODEL = "HuggingFaceTB/SmolVLM2-2.2B-Instruct",
LOCAL_MAX_NEW = 128,
MAX_IMAGE_SIDE = 1024,
JPEG_QUALITY = 90,
JUDGE = "rule",
NUM_REL_TOL = 0.0,
OUT_DIR = "/content/perceptionbench_out" if os.path.isdir("/content") else "./perceptionbench_out",
INSTALL_DEPS = True,
SHOW_PLOTS = True,
)
random.seed(CFG["SEED"])
os.makedirs(CFG["OUT_DIR"], exist_ok=True)
def _sh(pkgs):
subprocess.run([sys.executable, "-m", "pip", "install", "-q", *pkgs],
check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
if CFG["INSTALL_DEPS"]:
print("[setup] installing dependencies (quiet, ~30s on a cold Colab)…")
_sh(["datasets>=3.0.0", "huggingface_hub>=0.25.0", "pillow", "pandas",
"numpy", "matplotlib", "requests", "pyarrow"])
if CFG["BACKEND"] == "local":
_sh(["transformers>=4.51.0", "accelerate", "torch", "num2words"])
import numpy as np
import pandas as pd
import requests
import matplotlib
import matplotlib.pyplot as plt
from PIL import Image
matplotlib.rcParams.update({"figure.dpi": 110, "font.size": 9, "axes.grid": True,
"grid.alpha": .25, "axes.spines.top": False,
"axes.spines.right": False})
print("[setup] ready\n")
Loading the dataset
The code implements a resilient loader that attempts to stream converted Parquet files first. If that fails, it tries streaming the original data files. A full download occurs only as a last resort.
During the scan, the system limits the number of processed rows and organises examples into buckets based on the error_category field. This ensures a balanced sample where each visual capability contributes a comparable number of questions.
def _iter_rows(repo, split, mode, max_scan):
"""Yield dict rows, trying progressively heavier strategies."""
from datasets import load_dataset
if mode == "full":
print("[load] full download (~1.63 GB) …")
ds = load_dataset(repo, split=split)
for i, r in enumerate(ds):
if i >= max_scan:
return
yield r
return
try:
from huggingface_hub import HfApi, hf_hub_url
api = HfApi()
files = api.list_repo_files(repo, repo_type="dataset", revision="refs/convert/parquet")
pq = sorted(f for f in files if f.endswith(".parquet") and f"/{split}/" in f)
if pq:
urls = [hf_hub_url(repo, f, repo_type="dataset", revision="refs/convert/parquet") for f in pq]
print(f"[load] streaming {len(urls)} parquet shard(s) from refs/convert/parquet")
ds = load_dataset("parquet", data_files=urls, split="train", streaming=True)
for i, r in enumerate(ds):
if i >= max_scan:
return
yield r
return
except Exception as e:
print(f"[load] parquet stream unavailable ({type(e).__name__}: {e}); falling back")
try:
print("[load] streaming original data files")
ds = load_dataset(repo, split=split, streaming=True)
for i, r in enumerate(ds):
if i >= max_scan:
return
yield r
return
except Exception as e:
print(f"[load] json stream failed ({type(e).__name__}); doing a full download")
ds = load_dataset(repo, split=split)
for i, r in enumerate(ds):
if i >= max_scan:
return
yield r
def stratified_subset(repo, split, n_per_cat, max_scan, mode):
"""Balanced sample across `error_category` — the ten atomic capabilities.
Balancing matters: the benchmark reports a *capability profile*, and an
unbalanced sample makes the overall number a weighted average of whichever
capabilities happened to appear first in the shard.
"""
buckets, scanned, t0 = defaultdict(list), 0, time.time()
for row in _iter_rows(repo, split, mode, max_scan):
scanned += 1
cat = row.get("error_category") or "unknown"
if len(buckets[cat]) < n_per_cat:
buckets[cat].append(row)
if scanned % 100 == 0:
filled = sum(len(v) >= n_per_cat for v in buckets.values())
print(f" scanned={scanned:5d} categories={len(buckets):2d} "
f"filled={filled:2d} {time.time()-t0:5.1f}s", end="\r")
if scanned >= 250 and len(buckets) >= 10 and all(len(v) >= n_per_cat for v in buckets.values()):
break
rows = [r for v in buckets.values() for r in v]
random.Random(CFG["SEED"]).shuffle(rows)
print(f"\n[load] scanned {scanned} rows -> kept {len(rows)} across "
f"{len(buckets)} capabilities ({time.time()-t0:.1f}s)")
return rows, scanned
ROWS, N_SCANNED = stratified_subset(
CFG["REPO"], CFG["SPLIT"], CFG["N_PER_CATEGORY"], CFG["MAX_SCAN"], CFG["LOAD_MODE"])
Processing images and text
The workflow decodes base64-encoded images and parses interleaved image placeholders. It normalises each example into a consistent record format.
For performance, the system downscales images and re-encodes them as JPEGs with a quality setting of 90. This controls the token bill, as large screenshots can consume over 2,000 vision tokens per image.
The code also splits problem text on image placeholders to ensure no visual evidence is silently dropped. Any image not referenced by a placeholder is appended to the end of the list.
DATA_URI_RE = re.compile(r"^data:(image/[A-Za-z0-9.+-]+);base64,(.*)$", re.S)
PLACEHOLDER_RE = re.compile(r"<\|image[ _](\d+)\|>")
def decode_image(entry):
"""data-URI string | raw b64 | bytes | HF Image dict -> PIL.Image (RGB)."""
if isinstance(entry, Image.Image):
return entry.convert("RGB")
if isinstance(entry, dict):
if entry.get("bytes"):
return Image.open(io.BytesIO(entry["bytes"])).convert("RGB")
if entry.get("path"):
return Image.open(entry["path"]).convert("RGB")
if isinstance(entry, (bytes, bytearray)):
return Image.open(io.BytesIO(entry)).convert("RGB")
s = str(entry).strip()
m = DATA_URI_RE.match(s)
b64 = m.group(2) if m else s
b64 = re.sub(r"\s+", "", b64)
b64 += "=" * (-len(b64) % 4)
return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
def load_images(row):
imgs = row.get("image") or []
if isinstance(imgs, (str, bytes, dict)):
imgs = [imgs]
out = []
for e in imgs:
try:
out.append(decode_image(e))
except Exception as err:
print(f" [warn] undecodable image on idx={row.get('index')}: {err}")
return out
def shrink(img, max_side, quality):
"""Downscale + re-encode. Returns (PIL, data_uri). Controls the token bill:
a 3000px screenshot can cost >2k vision tokens per image, and these
questions carry up to 8 images each."""
w, h = img.size
if max(w, h) > max_side:
s = max_side / max(w, h)
img = img.resize((max(1, int(w * s)),Source Read original →Related articles
