| """The reward: how close a submitted document is to the ground truth. |
| |
| Two signals, because each is blind where the other sees: |
| |
| - **Structure** matches submitted elements to ground-truth elements one-to-one |
| (Hungarian assignment) and scores every pair on content, geometry and, for |
| text, typography. It is dense, it is readable off a trajectory, and it tells |
| you *which* element is wrong. Unmatched elements on either side count against |
| the score, so spamming elements does not buy coverage. |
| - **Visual** compares `render(submission)` with `render(ground_truth)` through |
| the same renderer: multi-scale SSIM plus colour-layout error, weighted toward |
| coarse scales and normalised so that an empty canvas scores zero rather than |
| whatever fraction of the target happens to be white. |
| |
| Reference mode (the model sees the target image) uses tight tolerances. |
| Description mode (the model only reads a description) loosens geometry, since |
| "top left, large" cannot pin a box to the pixel, and adds a coverage term: were |
| all the assets and all the words actually used? |
| |
| Every constant below is a tolerance a person can reason about: `pos_sigma` is a |
| fraction of the canvas diagonal, `size_sigma` a log-ratio. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import asdict, dataclass, field, replace |
|
|
| import numpy as np |
| from scipy.ndimage import gaussian_filter |
| from scipy.optimize import linear_sum_assignment |
|
|
| from .colors import delta_e |
| from .doc import DesignDoc, Element |
|
|
|
|
| @dataclass(frozen=True) |
| class ScoreConfig: |
| mode: str |
| pos_sigma: float |
| size_sigma: float |
| pixel_size: int |
| w_structure: float |
| w_visual: float |
| w_coverage: float |
| w_order: float |
| w_background: float |
| |
| |
| visual_terms: tuple = () |
|
|
|
|
| |
| |
| |
| REFERENCE = ScoreConfig("reference", 0.025, 0.12, 256, 0.60, 0.30, 0.00, 0.05, 0.05, |
| ((1, "ssim", 0.15), (2, "ssim", 0.30), (3, "ssim", 0.25), |
| (3, "color", 0.30))) |
| DESCRIPTION = ScoreConfig("description", 0.12, 0.45, 256, 0.50, 0.15, 0.25, 0.05, 0.05, |
| ((3, "ssim", 0.30), (4, "ssim", 0.30), (4, "color", 0.40))) |
| CONFIGS = {"reference": REFERENCE, "description": DESCRIPTION} |
| IOU_PAD = 0.004 |
|
|
| |
| |
| |
| |
| V3_COARSE_MIX = 0.5 |
| V3_COARSE_POS = 0.15 |
| V3_COARSE_SIZE = 0.6 |
| V3_TEXT_WEIGHT = 3.0 |
| V3_GATE = 0.3 |
|
|
|
|
| @dataclass |
| class PairScore: |
| gt_id: str |
| matched_id: str | None |
| score: float |
| content: float = 0.0 |
| geometry: float = 0.0 |
| style: float | None = None |
| note: str = "" |
|
|
|
|
| @dataclass |
| class Grade: |
| total: float |
| structure: float |
| visual: float |
| coverage: float |
| order: float |
| background: float |
| mode: str |
| n_gt: int |
| n_submitted: int |
| n_matched: int |
| elements: list[PairScore] = field(default_factory=list) |
| unmatched: list[str] = field(default_factory=list) |
|
|
| def as_dict(self) -> dict: |
| d = asdict(self) |
| for key in ("total", "structure", "visual", "coverage", "order", "background"): |
| d[key] = round(d[key], 4) |
| for e in d["elements"]: |
| for k in ("score", "content", "geometry"): |
| e[k] = round(e[k], 4) |
| if e["style"] is not None: |
| e["style"] = round(e["style"], 4) |
| return d |
|
|
| def report(self) -> str: |
| head = ( |
| f"reward {self.total:.3f} = structure {self.structure:.3f}, visual {self.visual:.3f}" |
| + (f", coverage {self.coverage:.3f}" if self.mode == "description" else "") |
| + f", order {self.order:.3f}, background {self.background:.3f}" |
| ) |
| lines = [head, f"matched {self.n_matched}/{self.n_gt} target elements, " |
| f"{self.n_submitted} submitted"] |
| worst = sorted(self.elements, key=lambda p: p.score)[:4] |
| for p in worst: |
| if p.score < 0.9: |
| lines.append(f" {p.gt_id}: {p.score:.2f} {p.note}".rstrip()) |
| if self.unmatched: |
| lines.append(f" extra elements: {', '.join(self.unmatched)}") |
| return "\n".join(lines) |
|
|
|
|
| |
|
|
|
|
| def levenshtein(a: str, b: str) -> int: |
| if a == b: |
| return 0 |
| if len(a) < len(b): |
| a, b = b, a |
| prev = list(range(len(b) + 1)) |
| for i, ca in enumerate(a, 1): |
| cur = [i] |
| for j, cb in enumerate(b, 1): |
| cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb))) |
| prev = cur |
| return prev[-1] |
|
|
|
|
| def text_similarity(a: str, b: str) -> float: |
| na = " ".join(a.split()).casefold() |
| nb = " ".join(b.split()).casefold() |
| if not na and not nb: |
| return 1.0 |
| return max(0.0, 1.0 - levenshtein(na, nb) / max(len(na), len(nb))) |
|
|
|
|
| def _box(el: Element) -> tuple[float, float, float, float]: |
| return el.x, el.y, el.x + el.width, el.y + el.height |
|
|
|
|
| def iou(a: Element, b: Element) -> float: |
| ax0, ay0, ax1, ay1 = _box(a) |
| bx0, by0, bx1, by1 = _box(b) |
| iw = max(0.0, min(ax1, bx1) - max(ax0, bx0)) |
| ih = max(0.0, min(ay1, by1) - max(ay0, by0)) |
| inter = iw * ih |
| union = a.width * a.height + b.width * b.height - inter |
| return inter / union if union > 0 else 0.0 |
|
|
|
|
| def _angle_gap(a: float, b: float) -> float: |
| d = abs(a - b) % 360 |
| return min(d, 360 - d) |
|
|
|
|
| def _iou_box(a, b) -> float: |
| iw = max(0.0, min(a[2], b[2]) - max(a[0], b[0])) |
| ih = max(0.0, min(a[3], b[3]) - max(a[1], b[1])) |
| inter = iw * ih |
| union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter |
| return inter / union if union > 0 else 0.0 |
|
|
|
|
| def geometry(a: Element, g: Element, diag: float, cfg: ScoreConfig, boxes=None) -> float: |
| """Position, size and overlap of what is drawn. |
| |
| `boxes` maps element ids to the box that is actually rendered; for text |
| that is the extent of the glyphs, not the declared box, whose width is |
| invisible for left-aligned copy and should not be scored. |
| """ |
| ba = (boxes or {}).get(("a", a.id)) or (a.x, a.y, a.x + a.width, a.y + a.height) |
| bg = (boxes or {}).get(("g", g.id)) or (g.x, g.y, g.x + g.width, g.y + g.height) |
| wa, ha = max(1e-3, ba[2] - ba[0]), max(1e-3, ba[3] - ba[1]) |
| wg, hg = max(1e-3, bg[2] - bg[0]), max(1e-3, bg[3] - bg[1]) |
| dist = math.dist(((ba[0] + ba[2]) / 2, (ba[1] + ba[3]) / 2), ((bg[0] + bg[2]) / 2, (bg[1] + bg[3]) / 2)) |
| pos = math.exp(-dist / (cfg.pos_sigma * diag)) |
| size = math.exp(-(abs(math.log(wa / wg)) + abs(math.log(ha / hg))) / (2 * cfg.size_sigma)) |
| rot = math.exp(-_angle_gap(a.angle, g.angle) / 10.0) |
| |
| |
| pad = IOU_PAD * diag |
| pa = (ba[0] - pad, ba[1] - pad, ba[2] + pad, ba[3] + pad) |
| pg = (bg[0] - pad, bg[1] - pad, bg[2] + pad, bg[3] + pad) |
| base = (_iou_box(pa, pg) + pos + size) / 3 |
| return base * (0.85 + 0.15 * rot) |
|
|
|
|
| def geometry_v3(a: Element, g: Element, diag: float, cfg: ScoreConfig, boxes=None) -> float: |
| """Coarse-to-fine geometry, scaled by the square root of the (coarse-to-fine) position match: an element in |
| the wrong place is the wrong element however well it is sized.""" |
| m = V3_COARSE_MIX |
| coarse = replace(cfg, pos_sigma=V3_COARSE_POS, size_sigma=V3_COARSE_SIZE) |
| base = (1 - m) * geometry(a, g, diag, cfg, boxes) + m * geometry(a, g, diag, coarse, boxes) |
| ba = (boxes or {}).get(("a", a.id)) or (a.x, a.y, a.x + a.width, a.y + a.height) |
| bg = (boxes or {}).get(("g", g.id)) or (g.x, g.y, g.x + g.width, g.y + g.height) |
| dist = math.dist(((ba[0] + ba[2]) / 2, (ba[1] + ba[3]) / 2), ((bg[0] + bg[2]) / 2, (bg[1] + bg[3]) / 2)) |
| pos = (1 - m) * math.exp(-dist / (cfg.pos_sigma * diag)) + m * math.exp(-dist / (V3_COARSE_POS * diag)) |
| return base * math.sqrt(pos) |
|
|
|
|
| def text_box(el: Element, fonts) -> tuple[float, float, float, float]: |
| from .text import layout |
|
|
| lay = layout(el, fonts) |
| x0 = min((ln.x for ln in lay.lines if ln.width > 0), default=0.0) |
| x1 = max((ln.x + ln.width for ln in lay.lines if ln.width > 0), default=el.width) |
| return el.x + x0, el.y, el.x + max(x1, x0 + 1.0), el.y + lay.height |
|
|
|
|
| def _dominant_color(el: Element) -> str: |
| runs = el.styled_runs() |
| weights: dict[str, int] = {} |
| for r in runs: |
| weights[r.color] = weights.get(r.color, 0) + len(r.text.strip()) |
| return max(weights.items(), key=lambda kv: kv[1])[0] if weights else (el.color or "#000000") |
|
|
|
|
| def typography(a: Element, g: Element, salient: bool = False) -> tuple[float, list[str]]: |
| notes = [] |
| font = 1.0 if (a.font_family or "").lower() == (g.font_family or "").lower() else 0.0 |
| if font < 1: |
| notes.append(f"font {a.font_family} vs {g.font_family}") |
| size = math.exp(-abs(math.log(a.font_size / g.font_size)) / 0.1) |
| if size < 0.7: |
| notes.append(f"size {a.font_size:.0f} vs {g.font_size:.0f}") |
| de = delta_e(_dominant_color(a), _dominant_color(g)) |
| color = math.exp(-de / 25.0) |
| if color < 0.7: |
| notes.append(f"color {_dominant_color(a)} vs {_dominant_color(g)}") |
| align = 1.0 if a.align == g.align else 0.0 |
| weight = 1.0 - 0.5 * ((a.bold != g.bold) + (a.italic != g.italic)) |
| case = 1.0 if a.uppercase == g.uppercase or a.display_text() == g.display_text() else 0.5 |
| spacing = math.exp(-abs(a.letter_spacing - g.letter_spacing) / max(1.0, 0.05 * g.font_size)) |
| lines = 1.0 |
| if len(g.lines) > 1: |
| lines = math.exp(-abs(a.line_height - g.line_height) / 0.15) |
| rest = (0.22 * size + 0.08 * align + 0.06 * weight + 0.04 * case + 0.05 * spacing + 0.05 * lines) / 0.50 |
| if salient: |
| return (0.5 + 0.5 * font) * (0.4 + 0.6 * color) * (0.5 + 0.5 * rest), notes |
| return 0.28 * font + 0.22 * color + 0.50 * rest, notes |
|
|
|
|
| def pair_score(a: Element, g: Element, diag: float, cfg: ScoreConfig, boxes=None, |
| asset_sim=None, style_weight: float = 0.38, version: int = 1) -> PairScore: |
| geometry_fn = geometry_v3 if version >= 3 else geometry |
| if a.kind == "text" and g.kind == "text": |
| content = text_similarity(a.display_text(), g.display_text()) |
| if content < 0.25: |
| return PairScore(g.id, a.id, 0.0, content, 0.0, None, "text differs") |
| geo = geometry_fn(a, g, diag, cfg, boxes) |
| style, notes = typography(a, g, salient=version >= 3) |
| opacity = math.exp(-abs(a.opacity - g.opacity) / 0.2) |
| if version >= 3: |
| score = content * geo * ((1 - style_weight) + style_weight * style) * (0.9 + 0.1 * opacity) |
| else: |
| score = content * ((1 - style_weight) * geo + style_weight * style) * (0.9 + 0.1 * opacity) |
| if content < 0.95: |
| notes.insert(0, f'text "{a.text[:20]}" vs "{g.text[:20]}"') |
| return PairScore(g.id, a.id, score, content, geo, style, "; ".join(notes)) |
| if a.kind == "text" or g.kind == "text": |
| return PairScore(g.id, a.id, 0.0, 0.0, 0.0, None, "kind mismatch") |
| if a.kind == "image" and g.kind == "image": |
| if a.asset_id == g.asset_id: |
| content = 1.0 |
| else: |
| |
| |
| sim = asset_sim(a.asset_id, g.asset_id) if asset_sim else 0.0 |
| content = 0.15 + 0.6 * sim |
| else: |
| content = 0.2 |
| geo = geometry_fn(a, g, diag, cfg, boxes) |
| opacity = math.exp(-abs(a.opacity - g.opacity) / 0.2) |
| note = "" if content == 1.0 else f"asset {a.asset_id or a.shape} vs {g.asset_id}" |
| return PairScore(g.id, a.id, content * geo * (0.9 + 0.1 * opacity), content, geo, None, note) |
|
|
|
|
| |
|
|
|
|
| def _gray(rgb: np.ndarray) -> np.ndarray: |
| return rgb[..., 0] * 0.299 + rgb[..., 1] * 0.587 + rgb[..., 2] * 0.114 |
|
|
|
|
| def ssim(a: np.ndarray, b: np.ndarray, sigma: float = 1.5) -> float: |
| x, y = _gray(a.astype(np.float64)), _gray(b.astype(np.float64)) |
| c1, c2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2 |
| mx, my = gaussian_filter(x, sigma), gaussian_filter(y, sigma) |
| sxx = gaussian_filter(x * x, sigma) - mx * mx |
| syy = gaussian_filter(y * y, sigma) - my * my |
| sxy = gaussian_filter(x * y, sigma) - mx * my |
| num = (2 * mx * my + c1) * (2 * sxy + c2) |
| den = (mx * mx + my * my + c1) * (sxx + syy + c2) |
| return float(np.clip((num / den).mean(), 0.0, 1.0)) |
|
|
|
|
| def pyramid(img: np.ndarray, levels: int) -> list[np.ndarray]: |
| """Area-averaged pyramid. Rendering small directly aliases (a one-pixel |
| canvas shift can flip a 32 px thumbnail's pixels by 150 levels); halving a |
| 256 px render does not.""" |
| out = [img.astype(np.float64)] |
| for _ in range(levels): |
| x = out[-1] |
| h, w = (x.shape[0] // 2) * 2, (x.shape[1] // 2) * 2 |
| if h < 4 or w < 4: |
| break |
| x = x[:h, :w] |
| out.append((x[0::2, 0::2] + x[1::2, 0::2] + x[0::2, 1::2] + x[1::2, 1::2]) / 4) |
| return out |
|
|
|
|
| def _term(kind: str, a: np.ndarray, b: np.ndarray) -> float: |
| if kind == "ssim": |
| return ssim(a, b) |
| mae = float(np.abs(a - b).mean() / 255.0) |
| return math.exp(-mae / 0.12) |
|
|
|
|
| def visual_score(attempt: np.ndarray, target: np.ndarray, blank: np.ndarray, |
| terms=REFERENCE.visual_terms) -> float: |
| """Multi-scale visual similarity, rescaled so an empty canvas scores 0. |
| |
| Each term is normalised against the blank canvas on its own scale, then |
| the terms are averaged with their weights. |
| """ |
| if attempt.shape != target.shape: |
| raise ValueError("renders differ in size") |
| depth = max(level for level, _, _ in terms) |
| pa, pt, pb = pyramid(attempt, depth), pyramid(target, depth), pyramid(blank, depth) |
| total = weight = 0.0 |
| for level, kind, w in terms: |
| level = min(level, len(pt) - 1) |
| v = _term(kind, pa[level], pt[level]) |
| v0 = _term(kind, pb[level], pt[level]) |
| score = (1.0 if v >= 0.999 else 0.0) if v0 >= 0.999 else (v - v0) / (1.0 - v0) |
| total += w * float(np.clip(score, 0.0, 1.0)) |
| weight += w |
| return total / weight if weight else 0.0 |
|
|
|
|
| |
|
|
|
|
| def _overlap(a: Element, b: Element) -> bool: |
| ax0, ay0, ax1, ay1 = _box(a) |
| bx0, by0, bx1, by1 = _box(b) |
| return min(ax1, bx1) > max(ax0, bx0) and min(ay1, by1) > max(ay0, by0) |
|
|
|
|
| def _split_runs(el: Element) -> list[list]: |
| """The element's runs cut at its line breaks: one run list per line.""" |
| per_line: list[list] = [[]] |
| for r in el.runs or []: |
| parts = r.text.split("\n") |
| for k, part in enumerate(parts): |
| if k: |
| per_line.append([]) |
| if part: |
| per_line[-1].append(r.model_copy(update={"text": part})) |
| return per_line |
|
|
|
|
| def explode_lines(elements: list[Element], fonts) -> list[Element]: |
| """Reward v2 scores text line by line: a three-line box and three stacked one-line boxes that |
| render identically are the same design, and building it either way should score the same.""" |
| from .text import layout |
|
|
| out = [] |
| for e in elements: |
| if e.kind != "text" or e.angle or len(e.lines) < 2 or fonts is None: |
| out.append(e) |
| continue |
| lay = layout(e, fonts) |
| runs = _split_runs(e) if e.runs else None |
| for k, (line, text) in enumerate(zip(lay.lines, e.lines, strict=False)): |
| if not text.strip(): |
| continue |
| out.append(e.model_copy(update={ |
| "id": f"{e.id}#l{k}", "text": text, "y": e.y + line.top, "height": lay.line_box, |
| "runs": (runs[k] or None) if runs and k < len(runs) else None})) |
| return out |
|
|
|
|
| def grade( |
| attempt: DesignDoc, |
| target: DesignDoc, |
| mode: str = "reference", |
| attempt_pixels: np.ndarray | None = None, |
| target_pixels: np.ndarray | None = None, |
| blank_pixels: np.ndarray | None = None, |
| fonts=None, |
| asset_sim=None, |
| version: int = 1, |
| ) -> Grade: |
| cfg = CONFIGS[mode] |
| if version >= 3 and mode != "reference": |
| version = 1 |
| sub = [e for e in attempt.elements if e.visible and e.opacity > 0.01] |
| gts = [e for e in target.elements if e.visible and e.opacity > 0.01] |
| if version == 2: |
| sub, gts = explode_lines(sub, fonts), explode_lines(gts, fonts) |
| diag = math.hypot(target.width, target.height) |
| boxes = {} |
| if fonts is not None: |
| |
| |
| boxes.update({("a", e.id): text_box(e, fonts) for e in sub if e.kind == "text" and not e.angle}) |
| boxes.update({("g", e.id): text_box(e, fonts) for e in gts if e.kind == "text" and not e.angle}) |
|
|
| if not sub: |
| return Grade(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, mode, len(gts), 0, 0, |
| [PairScore(g.id, None, 0.0, note="missing") for g in gts], []) |
|
|
| matrix = np.zeros((len(sub), len(gts))) |
| pairs: dict[tuple[int, int], PairScore] = {} |
| for i, a in enumerate(sub): |
| for j, g in enumerate(gts): |
| |
| |
| ps = pair_score(a, g, diag, cfg, boxes, asset_sim, style_weight=0.5 if version == 2 else 0.38, |
| version=version) |
| matrix[i, j] = ps.score |
| pairs[(i, j)] = ps |
| rows, cols = linear_sum_assignment(-matrix) if gts else ([], []) |
| matched = {} |
| for i, j in zip(rows, cols, strict=True): |
| if matrix[i, j] > 0: |
| matched[j] = i |
| elements = [] |
| for j, g in enumerate(gts): |
| if j in matched: |
| elements.append(pairs[(matched[j], j)]) |
| else: |
| elements.append(PairScore(g.id, None, 0.0, note="missing")) |
| used = set(matched.values()) |
| unmatched = [a.id for i, a in enumerate(sub) if i not in used] |
| if version >= 3: |
| |
| weights = [V3_TEXT_WEIGHT if g.kind == "text" else 1.0 for g in gts] |
| mean_w = sum(weights) / len(weights) if weights else 1.0 |
| denom = sum(weights) + len(unmatched) * mean_w |
| structure = sum(w * p.score for w, p in zip(weights, elements, strict=True)) / denom if denom else 1.0 |
| else: |
| denom = max(len(gts), len(sub)) |
| structure = sum(p.score for p in elements) / denom if denom else 1.0 |
|
|
| |
| agree = total = 0 |
| gt_index = {g.id: k for k, g in enumerate(gts)} |
| sub_index = {a.id: k for k, a in enumerate(sub)} |
| matched_pairs = [(gts[j], sub[i]) for j, i in matched.items()] |
| for p in range(len(matched_pairs)): |
| for q in range(p + 1, len(matched_pairs)): |
| g1, a1 = matched_pairs[p] |
| g2, a2 = matched_pairs[q] |
| if not _overlap(g1, g2): |
| continue |
| total += 1 |
| same = (gt_index[g1.id] < gt_index[g2.id]) == (sub_index[a1.id] < sub_index[a2.id]) |
| agree += same |
| order = (agree / total if total else 1.0) if matched else 0.0 |
|
|
| background = math.exp(-delta_e(attempt.background, target.background) / 20.0) |
|
|
| visual = 0.0 |
| if attempt_pixels is not None and target_pixels is not None and blank_pixels is not None: |
| visual = visual_score(attempt_pixels, target_pixels, blank_pixels, cfg.visual_terms) |
|
|
| coverage = 0.0 |
| if mode == "description": |
| gt_assets = {g.asset_id for g in gts if g.kind == "image"} |
| sub_assets = {a.asset_id for a in sub if a.kind == "image"} |
| asset_cov = len(gt_assets & sub_assets) / len(gt_assets) if gt_assets else 1.0 |
| gt_texts = [g.display_text() for g in gts if g.kind == "text"] |
| sub_texts = [a.display_text() for a in sub if a.kind == "text"] |
| hits = sum( |
| 1 for t in gt_texts if any(text_similarity(t, s) >= 0.8 for s in sub_texts) |
| ) |
| text_cov = hits / len(gt_texts) if gt_texts else 1.0 |
| coverage = 0.5 * asset_cov + 0.5 * text_cov |
|
|
| gate = min(1.0, structure / V3_GATE) if version >= 3 else 1.0 |
| total_score = ( |
| cfg.w_structure * structure |
| + gate * (cfg.w_visual * visual + cfg.w_order * order + cfg.w_background * background) |
| + cfg.w_coverage * coverage |
| ) |
| return Grade( |
| total=float(min(1.0, max(0.0, total_score))), |
| structure=float(structure), |
| visual=float(visual), |
| coverage=float(coverage), |
| order=float(order), |
| background=float(background), |
| mode=mode, |
| n_gt=len(gts), |
| n_submitted=len(sub), |
| n_matched=len(matched), |
| elements=elements, |
| unmatched=unmatched, |
| ) |
|
|
