| """The task pack: prepared tasks on disk, and the catalog the server reads them from. |
| |
| Layout: |
| |
| <pack>/ |
| manifest.json version, source revision, split counts |
| fonts/manifest.json + files the 263 Crello families, extracted |
| tasks/<split>.jsonl one task per line (ground truth included) |
| designs/<crello_id>/ assets/<aid>.png|webp and preview.jpg |
| |
| A task is one Crello template. How it is played is chosen at reset: `mode` |
| (reference image or text description) and `surface` (tools, HTML or GUI) are |
| episode parameters, not separate tasks, so every split serves every variant and |
| the Task API stays a plain list of designs. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| import threading |
| import time |
| from collections import OrderedDict |
| from pathlib import Path |
|
|
| from .describe import describe |
| from .doc import DesignDoc |
| from .fonts import FontRegistry, extract_crello_fonts |
| from .render import Renderer |
|
|
| PACK_VERSION = 1 |
| SPLITS = ("train", "validation", "test") |
| DEFAULT_PACK = Path(os.environ.get("DESIGN_CANVAS_PACK", Path.home() / ".cache/design_canvas/pack")) |
| CRELLO_REPO = "cyberagent/crello" |
| CRELLO_REVISION = "5.1.0" |
|
|
|
|
| def difficulty(n: int) -> str: |
| return "easy" if n <= 6 else "medium" if n <= 12 else "hard" |
|
|
|
|
| class TaskPack: |
| """Read-only view of a prepared pack. Safe to share across sessions and threads.""" |
|
|
| def __init__(self, root: Path | str): |
| self.root = Path(root) |
| manifest = self.root / "manifest.json" |
| if not manifest.exists(): |
| raise FileNotFoundError( |
| f"no task pack at {self.root}. Build one with `design-canvas prepare` " |
| f"or point DESIGN_CANVAS_PACK at an existing pack." |
| ) |
| self.manifest = json.loads(manifest.read_text()) |
| self.fonts = FontRegistry(self.root / "fonts") |
| self.renderer = Renderer(self.fonts) |
| self._tasks: dict[str, list[dict]] = {} |
| self._by_id: dict[str, dict] = {} |
| self._docs: dict[str, DesignDoc] = {} |
| self._renders: OrderedDict[tuple, object] = OrderedDict() |
| self._lock = threading.Lock() |
| for split in self.splits(): |
| path = self.root / "tasks" / f"{split}.jsonl" |
| tasks = [json.loads(line) for line in path.read_text().splitlines() if line.strip()] |
| self._tasks[split] = tasks |
| for t in tasks: |
| self._by_id[t["task_id"]] = t |
|
|
| |
|
|
| def splits(self) -> list[str]: |
| return [s for s in SPLITS if s in self.manifest.get("splits", {})] |
|
|
| def count(self, split: str) -> int: |
| if split not in self._tasks: |
| raise KeyError(f"unknown split {split!r}; have {', '.join(self.splits())}") |
| return len(self._tasks[split]) |
|
|
| def at(self, split: str, index: int) -> dict: |
| tasks = self._tasks.get(split) |
| if tasks is None: |
| raise KeyError(f"unknown split {split!r}; have {', '.join(self.splits())}") |
| if not 0 <= index < len(tasks): |
| raise IndexError(f"{split} has {len(tasks)} tasks; index {index} is out of range") |
| return tasks[index] |
|
|
| def get(self, task_id: str) -> dict: |
| try: |
| return self._by_id[task_id] |
| except KeyError: |
| raise KeyError(f"unknown task_id {task_id!r}") from None |
|
|
| @staticmethod |
| def public(task: dict) -> dict: |
| """The Task API view: enough to choose a task, nothing that gives it away.""" |
| return { |
| "task_id": task["task_id"], |
| "index": task["index"], |
| "split": task["split"], |
| "crello_id": task["crello_id"], |
| "title": task["title"], |
| "format": task["format"], |
| "category": task["category"], |
| "canvas": task["canvas"], |
| "n_elements": task["stats"]["n_elements"], |
| "n_text": task["stats"]["n_text"], |
| "n_assets": len(task["assets"]), |
| "difficulty": task["difficulty"], |
| "modes": ["reference", "description"], |
| "surfaces": ["tools", "html", "gui"], |
| } |
|
|
| def task_range(self, split: str, start=None, stop=None) -> list[dict]: |
| tasks = self._tasks.get(split) |
| if tasks is None: |
| raise KeyError(f"unknown split {split!r}") |
| return [self.public(t) for t in tasks[slice(start, stop)]] |
|
|
| |
|
|
| def gt_doc(self, task: dict) -> DesignDoc: |
| tid = task["task_id"] |
| with self._lock: |
| doc = self._docs.get(tid) |
| if doc is None: |
| doc = DesignDoc.model_validate(task["doc"]) |
| self._docs[tid] = doc |
| return doc |
|
|
| def asset_paths(self, task: dict) -> dict[str, Path]: |
| return {a["id"]: self.root / a["file"] for a in task["assets"]} |
|
|
| def asset_ctx(self, task: dict) -> dict[str, dict]: |
| return {a["id"]: {"aspect": a["aspect"], "kind": a["kind"]} for a in task["assets"]} |
|
|
| def preview_path(self, task: dict) -> Path: |
| return self.root / task["preview"] |
|
|
| def _cached(self, key: tuple, make): |
| with self._lock: |
| if key in self._renders: |
| self._renders.move_to_end(key) |
| return self._renders[key] |
| value = make() |
| with self._lock: |
| self._renders[key] = value |
| while len(self._renders) > 256: |
| self._renders.popitem(last=False) |
| return value |
|
|
| def target_png(self, task: dict, max_size: int = 512) -> bytes: |
| return self._cached( |
| ("png", task["task_id"], max_size), |
| lambda: self.renderer.render_png(self.gt_doc(task), self.asset_paths(task), max_size), |
| ) |
|
|
| def target_jpeg(self, task: dict, max_size: int = 768) -> bytes: |
| return self._cached( |
| ("jpg", task["task_id"], max_size), |
| lambda: self.renderer.render_jpeg(self.gt_doc(task), self.asset_paths(task), max_size), |
| ) |
|
|
| def target_pixels(self, task: dict, max_size: int): |
| return self._cached( |
| ("arr", task["task_id"], max_size), |
| lambda: self.renderer.render_array(self.gt_doc(task), self.asset_paths(task), max_size), |
| ) |
|
|
| def asset_similarity(self, task: dict): |
| """`sim(a, b)` in [0, 1]: how alike two of this task's assets look. |
| |
| Compared on 32 px thumbnails composited over grey (so the cut-out shape |
| counts), on average colour, and on aspect ratio. Squared, so only |
| genuine look-alikes earn real credit. |
| """ |
|
|
| def make(): |
| import numpy as np |
| from PIL import Image |
|
|
| from .score import ssim |
|
|
| thumbs, colors, aspect = {}, {}, {} |
| for a in task["assets"]: |
| with Image.open(self.root / a["file"]) as im: |
| im = im.convert("RGBA") |
| bg = Image.new("RGBA", im.size, (128, 128, 128, 255)) |
| flat = Image.alpha_composite(bg, im).convert("RGB") |
| thumbs[a["id"]] = np.asarray(flat.resize((32, 32), Image.BOX), dtype=np.float64) |
| colors[a["id"]] = np.asarray(flat.resize((4, 4), Image.BOX), dtype=np.float64) |
| aspect[a["id"]] = a["aspect"] |
| table = {} |
| ids = list(thumbs) |
| for i, x in enumerate(ids): |
| for y in ids[i + 1:]: |
| shape = ssim(thumbs[x], thumbs[y]) |
| color = float(np.exp(-np.abs(colors[x] - colors[y]).mean() / 255 / 0.1)) |
| ratio = float(np.exp(-abs(np.log(aspect[x] / aspect[y])) / 0.25)) |
| sim = (0.4 * shape + 0.3 * color + 0.3 * ratio) ** 2 |
| table[(x, y)] = table[(y, x)] = round(sim, 4) |
| return table |
|
|
| table = self._cached(("assetsim", task["task_id"]), make) |
| return lambda a, b: 1.0 if a == b else table.get((a, b), 0.0) |
|
|
| def grade(self, task: dict, doc: DesignDoc, mode: str = "reference", version: int | None = None): |
| """The one scoring entry point: episodes, tests and scripts all call this.""" |
| from .score import CONFIGS, grade |
|
|
| size = CONFIGS[mode].pixel_size |
| attempt = self.renderer.render_array(doc, self.asset_paths(task), size) |
| return grade(doc, self.gt_doc(task), mode, attempt, self.target_pixels(task, size), |
| self.blank_pixels(task, size), fonts=self.fonts, |
| asset_sim=self.asset_similarity(task), |
| version=version or int(os.environ.get("DESIGN_CANVAS_REWARD", "1"))) |
|
|
| def blank_pixels(self, task: dict, max_size: int): |
| def make(): |
| doc = self.gt_doc(task) |
| return self.renderer.render_array( |
| DesignDoc(width=doc.width, height=doc.height), {}, max_size |
| ) |
|
|
| return self._cached(("blank", task["task_id"], max_size), make) |
|
|
|
|
| _PACKS: dict[str, TaskPack] = {} |
| _PACKS_LOCK = threading.Lock() |
|
|
|
|
| def load_pack(root: Path | str | None = None) -> TaskPack: |
| """The task catalog the server plays from. |
| |
| `DESIGN_CANVAS_CORPUS` (a corpus manifest: local path or `hf://buckets/...`) |
| selects the full bucket-backed corpus; otherwise a prepared pack directory |
| (`root`, `DESIGN_CANVAS_PACK`, or the default) is served. |
| """ |
| corpus = None if root else os.environ.get("DESIGN_CANVAS_CORPUS") |
| key = corpus or str(Path(root or DEFAULT_PACK).expanduser().resolve()) |
| with _PACKS_LOCK: |
| pack = _PACKS.get(key) |
| if pack is None: |
| if corpus: |
| from .corpus import CorpusPack |
|
|
| pack = CorpusPack(corpus) |
| else: |
| pack = TaskPack(key) |
| _PACKS[key] = pack |
| return pack |
|
|
|
|
| |
|
|
|
|
| def accept(design, max_elements: int) -> str | None: |
| """Why a design is rejected, or None to keep it.""" |
| n = len(design.doc.elements) |
| if n < 2: |
| return "fewer than two elements" |
| if n > max_elements: |
| return f"more than {max_elements} elements" |
| if not any(e.kind == "text" for e in design.doc.elements): |
| return "no text" |
| if any(w.startswith("font fallback") for w in design.warnings): |
| return "font not in the Crello font set" |
| if max(design.doc.width, design.doc.height) > 4000: |
| return "canvas too large" |
| return None |
|
|
|
|
| def _asset_ext(kind: str) -> str: |
| """Photos as WebP (they are most of the bytes), graphics as PNG (crisp edges).""" |
| return ".webp" if kind in ("photo", "background") else ".png" |
|
|
|
|
| def write_design(out: Path, design) -> None: |
| """Write a converted design's assets and Crello preview under `out/designs/<id>/`.""" |
| ddir = Path(out) / "designs" / design.crello_id |
| (ddir / "assets").mkdir(parents=True, exist_ok=True) |
| for aid, img in design.assets.items(): |
| ext = _asset_ext(design.asset_meta[aid]["kind"]) |
| tmp = ddir / "assets" / f".{aid}{ext}.part" |
| if ext == ".webp": |
| img.save(tmp, "WEBP", quality=90, method=4) |
| else: |
| img.save(tmp, "PNG", optimize=True) |
| tmp.replace(ddir / "assets" / f"{aid}{ext}") |
| if design.preview is not None: |
| tmp = ddir / ".preview.jpg.part" |
| design.preview.convert("RGB").save(tmp, "JPEG", quality=88) |
| tmp.replace(ddir / "preview.jpg") |
|
|
|
|
| def make_task(design, split: str, index: int) -> dict: |
| """The task record for a converted design: ground truth, asset list, brief, stats.""" |
| cid = design.crello_id |
| assets = sorted( |
| ({**meta, "file": f"designs/{cid}/assets/{aid}{_asset_ext(meta['kind'])}"} |
| for aid, meta in design.asset_meta.items()), |
| key=lambda a: int(a["id"][1:]), |
| ) |
| n_text = sum(e.kind == "text" for e in design.doc.elements) |
| return { |
| "task_id": f"{split}-{index:05d}", |
| "index": index, |
| "split": split, |
| **design.meta, |
| "doc": design.doc.model_dump(mode="json", exclude_defaults=False), |
| "assets": assets, |
| "fonts": sorted({e.font_family for e in design.doc.elements if e.kind == "text"}), |
| "preview": f"designs/{cid}/preview.jpg", |
| "description": describe(design.doc, design.meta, design.asset_meta), |
| "stats": { |
| "n_elements": len(design.doc.elements), |
| "n_text": n_text, |
| "n_images": len(design.doc.elements) - n_text, |
| }, |
| "difficulty": difficulty(len(design.doc.elements)), |
| } |
|
|
|
|
| def build_pack( |
| out: Path | str, |
| sources: dict[str, list[str]], |
| per_split: dict[str, int], |
| fonts_pickle: str, |
| max_elements: int = 20, |
| log=print, |
| ) -> dict: |
| """Convert Crello parquet shards into a task pack. |
| |
| `sources` maps split -> local parquet paths; `per_split` caps how many |
| accepted designs each split keeps, taken in file order. |
| """ |
| import datasets |
| import pyarrow.parquet as pq |
|
|
| from .crello import convert_row |
|
|
| out = Path(out) |
| (out / "tasks").mkdir(parents=True, exist_ok=True) |
| if not (out / "fonts" / "manifest.json").exists(): |
| extract_crello_fonts(fonts_pickle, out / "fonts") |
| fonts = FontRegistry(out / "fonts") |
| features = datasets.load_dataset_builder(CRELLO_REPO, revision=CRELLO_REVISION).info.features |
|
|
| counts = {} |
| rejected: dict[str, int] = {} |
| for split, paths in sources.items(): |
| want = per_split.get(split, 0) |
| tasks = [] |
| for path in paths: |
| if len(tasks) >= want: |
| break |
| pf = pq.ParquetFile(path) |
| for rg in range(pf.num_row_groups): |
| if len(tasks) >= want: |
| break |
| for row in pf.read_row_group(rg).to_pylist(): |
| if len(tasks) >= want: |
| break |
| design = convert_row(row, features, fonts) |
| reason = accept(design, max_elements) |
| if reason: |
| rejected[reason] = rejected.get(reason, 0) + 1 |
| continue |
| write_design(out, design) |
| tasks.append(make_task(design, split, len(tasks))) |
| log(f"{split}: {len(tasks)}/{want} after {Path(path).name}") |
| with open(out / "tasks" / f"{split}.jsonl", "w") as fh: |
| for t in tasks: |
| fh.write(json.dumps(t, separators=(",", ":")) + "\n") |
| counts[split] = len(tasks) |
|
|
| manifest = { |
| "version": PACK_VERSION, |
| "created": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), |
| "source": {"repo": CRELLO_REPO, "revision": CRELLO_REVISION, |
| "license": "CDLA-Permissive-2.0; templates (c) VistaCreate, see dataset card"}, |
| "splits": counts, |
| "max_elements": max_elements, |
| "rejected": rejected, |
| } |
| (out / "manifest.json").write_text(json.dumps(manifest, indent=2)) |
| return manifest |
|
|
