| |
| """Element-head teacher data, step 1: contact sheets showing each page image every way the extension can render it. |
| |
| For each sampled image (logo, icon, formula, photo, backdrop...) of a crawled light page that has a Haiku dark design, |
| one row: the element in its light page (red outline), then the candidates on the dark surface designed for that spot: |
| A keep - unchanged |
| B invert - exactly Chrome's `invert(1) hue-rotate(180deg)` (dark <-> light, hues kept) |
| C adapt - the runtime's adaptPixels: neutral dark ink -> light, white box removed, brand colours kept |
| D remove - large CSS background images only: the designed dark surface shows through |
| Writes OUT/sheets/<id>.png and OUT/meta.<worker>.jsonl (one line per sheet: rows with page id, image index, offers). |
| """ |
| import argparse, base64, glob, gzip, hashlib, io, json, math, os, random, sys |
| from multiprocessing import Pool |
| from pathlib import Path |
|
|
| import numpy as np |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| HUE180 = np.array([[-0.574, 1.43, 0.144], [0.426, 0.43, 0.144], [0.426, 1.43, -0.856]]) |
| SMALL, LARGE = "small", "large" |
| GEOM = {SMALL: dict(cell=(124, 112), ctx=(232, 112), rows=8), LARGE: dict(cell=(250, 140), ctx=(250, 140), rows=5)} |
| KIND = {"img": "image", "bg": "CSS background", "pseudo": "CSS icon", "svg": "inline SVG", "svgimage": "SVG image"} |
|
|
|
|
| |
| def lin(u): |
| u = u / 255.0 |
| return np.where(u <= 0.04045, u / 12.92, ((u + 0.055) / 1.055) ** 2.4) |
|
|
|
|
| def oklab(rgb): |
| r, g, b = lin(rgb[..., 0]), lin(rgb[..., 1]), lin(rgb[..., 2]) |
| l = np.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b) |
| m = np.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b) |
| s = np.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b) |
| return np.stack([0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s, |
| 1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s, |
| 0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s], -1) |
|
|
|
|
| def to_rgb(L, a, b): |
| l_ = L + 0.3963377774 * a + 0.2158037573 * b |
| m_ = L - 0.1055613458 * a - 0.0638541728 * b |
| s_ = L - 0.0894841775 * a - 1.2914855480 * b |
| l, m, s = l_ ** 3, m_ ** 3, s_ ** 3 |
| def enc(x): |
| x = np.clip(x, 0, 1) |
| return np.round(255 * np.where(x <= 0.0031308, 12.92 * x, 1.055 * np.power(x, 1 / 2.4) - 0.055)) |
| return np.stack([enc(4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s), |
| enc(-1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s), |
| enc(-0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s)], -1) |
|
|
|
|
| def invert(rgba): |
| x = 1.0 - rgba[..., :3] / 255.0 |
| out = rgba.copy() |
| out[..., :3] = np.round(np.clip(x @ HUE180.T, 0, 1) * 255) |
| return out |
|
|
|
|
| def adapt(rgba): |
| """Port of nadir-map.js adaptPixels (incl. transparent mode). None when there is nothing to adapt.""" |
| d = rgba.astype(np.float64) |
| h, w = d.shape[:2] |
| ring = np.zeros((h, w), bool) |
| ring[0, :] = ring[-1, :] = ring[:, 0] = ring[:, -1] = True |
| rp = d[ring] |
| ringN = len(rp) |
| opq = rp[:, 3] >= 250 |
| n = int(opq.sum()) |
| o = oklab(rp[opq][:, :3]) if n else np.zeros((0, 3)) |
| lightm = (o[:, 0] > 0.9) & (np.hypot(o[:, 1], o[:, 2]) < 0.04) if n else np.zeros(0, bool) |
| light = int(lightm.sum()) |
| bg = None |
| if n and light / n >= 0.5: |
| bg = oklab(rp[opq][lightm][:, :3].mean(0)) |
| elif n > 0.2 * ringN: |
| return None |
| else: |
| a = d[..., 3] >= 128 |
| if not a.any(): |
| return None |
| ol = oklab(d[a][:, :3]) |
| ink = (ol[:, 0] < 0.5) & (np.hypot(ol[:, 1], ol[:, 2]) < 0.05) |
| if ink.mean() < 0.08: |
| return None |
| out = np.zeros_like(d) |
| vis = d[..., 3] > 0 |
| px = d[vis] |
| ol = oklab(px[:, :3]) |
| if bg is not None: |
| dist = np.linalg.norm(ol - bg, axis=1) |
| cov = np.clip((dist - 0.03) / 0.22, 0, 1) |
| else: |
| cov = np.ones(len(px)) |
| keep = cov > 0 |
| c = np.where(keep, cov, 1.0) |
| if bg is not None: |
| L = (ol[:, 0] - (1 - c) * bg[0]) / c |
| A = (ol[:, 1] - (1 - c) * bg[1]) / c |
| B = (ol[:, 2] - (1 - c) * bg[2]) / c |
| else: |
| L, A, B = ol[:, 0].copy(), ol[:, 1].copy(), ol[:, 2].copy() |
| L = np.clip(L, 0, 1) |
| C = np.hypot(A, B) |
| neutral = C < 0.05 |
| lift = (~neutral) & (L < 0.62) |
| L = np.where(neutral, 0.94 - 0.75 * L, np.where(lift, 0.62 + 0.15 * L / 0.62, L)) |
| A = np.where(lift, A * 0.92, A) |
| B = np.where(lift, B * 0.92, B) |
| rgb = to_rgb(L, A, B) |
| res = np.zeros((len(px), 4)) |
| res[:, :3] = rgb |
| res[:, 3] = np.round(px[:, 3] * cov) |
| res[~keep] = 0 |
| out[vis] = res |
| return out.astype(np.uint8) |
|
|
|
|
| |
| def hex_rgba(h): |
| h = h.lstrip("#") |
| r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16) |
| a = int(h[6:8], 16) / 255 if len(h) >= 8 else 1.0 |
| return r, g, b, a |
|
|
|
|
| def over(top, base): |
| r, g, b, a = top |
| return tuple(round(a * t + (1 - a) * u) for t, u in zip((r, g, b), base)) |
|
|
|
|
| def dark_palette(page, design): |
| """Dark surface per token index and the main dark text colour, from the Haiku design.""" |
| cols = {int(k): v for k, v in design["colors"].items()} if isinstance(design["colors"], dict) else {int(c["id"].lstrip("t")): c["dark"] for c in design["colors"]} |
| toks = page["tokens"] |
| root = next((i for i, t in enumerate(toks) if t["role"] == 0 and t["f"][24] > 0.5), 0) |
| canvas = hex_rgba(cols[root])[:3] if root in cols else (18, 18, 20) |
| best, text = -1, (230, 232, 235) |
| for i, t in enumerate(toks): |
| if t["role"] == 1 and i in cols and t["f"][19] > best: |
| best, text = t["f"][19], over(hex_rgba(cols[i]), canvas) |
| surf = {i: over(hex_rgba(cols[i]), canvas) for i, t in enumerate(toks) if i in cols and t["role"] == 0} |
| return surf, canvas, text |
|
|
|
|
| def content_area(o, big): |
| """The big thumbnail is letterboxed into a square: cut out the image itself.""" |
| S = big.shape[0] |
| if o.get("cropPx"): |
| nw, nh = o["cropPx"][2], o["cropPx"][3] |
| elif o.get("kind") == "svg": |
| nw, nh = o["w"], o["h"] |
| else: |
| nw, nh = (o.get("nw") or o["w"]), (o.get("nh") or o["h"]) |
| nw, nh = max(1, nw), max(1, nh) |
| s = min(S / nw, S / nh) |
| dw, dh = max(1, round(nw * s)), max(1, round(nh * s)) |
| x0, y0 = (S - dw) // 2, (S - dh) // 2 |
| return big[y0:y0 + dh, x0:x0 + dw] |
|
|
|
|
| def decode_png(data_url): |
| return np.array(Image.open(io.BytesIO(base64.b64decode(data_url.split(",", 1)[1]))).convert("RGBA")) |
|
|
|
|
| def bucket(o): |
| s = o.get("stats") or {} |
| area = o["w"] * o["h"] |
| if o.get("logo") or (o.get("homeLink") and o.get("header")): |
| return "logo" |
| if o["kind"] == "bg" and area >= 20000: |
| return "backdrop" |
| if o["kind"] in ("pseudo",) or o.get("cropPx"): |
| return "icon" |
| if s.get("transp", 0) > 0.15: |
| return "transparent" |
| if s.get("transp", 1) < 0.02 and s.get("colors", 0) > 0.18 and s.get("stdL", 0) > 0.12: |
| return "photo" |
| return "flat" |
|
|
|
|
| RATE = {"logo": 1.0, "backdrop": 1.0, "icon": 1.0, "transparent": 1.0, "flat": 0.8, "photo": 0.12} |
|
|
|
|
| def candidates(task): |
| """All sampled elements of one page with everything needed to draw their rows.""" |
| crawl, haiku, pid, tag, seed = task |
| rnd = random.Random(seed) |
| try: |
| page = json.load(gzip.open(Path(crawl) / "pages" / pid / "light.json.gz")) |
| design = json.load(open(Path(haiku) / f"{pid}.json")) |
| if design.get("refusal") or not design.get("colors"): |
| return [] |
| surf, canvas, text = dark_palette(page, design) |
| except Exception: |
| return [] |
| out = [] |
| for i, o in enumerate(page["images"]): |
| if not o.get("big") or o["w"] * o["h"] < 64 or o["w"] < 4 or o["h"] < 4: |
| continue |
| if o["kind"] == "svg" and not o.get("useExt"): |
| continue |
| b = bucket(o) |
| if rnd.random() > RATE[b]: |
| continue |
| dsurf = surf.get(o.get("underTok", -1), canvas) |
| out.append({"page": pid, "crawl": tag, "img": i, "bucket": b, "o": {k: o.get(k) for k in ("kind", "src", "x", "y", "w", "h", "nw", "nh", "cropPx", "logo", "homeLink", "header", "inLink", "surf", "big", "stats")}, |
| "dsurf": dsurf, "dtext": text, "light": tuple((o.get("surf") or [255, 255, 255])[:3])}) |
| return out |
|
|
|
|
| |
| def font(sz): |
| try: |
| return ImageFont.load_default(size=sz) |
| except TypeError: |
| return ImageFont.load_default() |
|
|
|
|
| F11, F13, F15 = font(11), font(13), font(15) |
|
|
|
|
| def fit(o, box, maxzoom=4.0): |
| s = min(box[0] / max(1, o["w"]), box[1] / max(1, o["h"]), maxzoom) |
| return max(1, round(o["w"] * s)), max(1, round(o["h"] * s)), s |
|
|
|
|
| def variant_cell(img, o, typ, size, surface, text, mode): |
| W, H = size |
| cell = Image.new("RGBA", (W, H), tuple(surface) + (255,)) |
| if mode != "D": |
| dw, dh, _ = fit(o, (W - 12, H - 12)) |
| im = Image.fromarray(img, "RGBA").resize((dw, dh), Image.LANCZOS) |
| cell.alpha_composite(im, ((W - dw) // 2, (H - dh) // 2)) |
| if typ == LARGE and o["kind"] == "bg": |
| dr = ImageDraw.Draw(cell) |
| dr.text((12, H // 2 - 22), "Heading on this background", fill=tuple(text), font=F15) |
| dr.text((12, H // 2 + 2), "Body text that sits on top of it.", fill=tuple(text), font=F13) |
| return cell |
|
|
|
|
| def context_cell(shot, e, size): |
| W, H = size |
| o = e["o"] |
| if shot is not None and o["y"] + o["h"] <= shot.height and o["x"] >= -2 and o["x"] + o["w"] <= shot.width + 2: |
| m = max(30, 0.75 * max(o["w"], o["h"])) |
| x0, y0 = max(0, o["x"] - m), max(0, o["y"] - m) |
| x1, y1 = min(shot.width, o["x"] + o["w"] + m), min(shot.height, o["y"] + o["h"] + m) |
| crop = shot.crop((int(x0), int(y0), int(x1), int(y1))) |
| s = min(W / crop.width, H / crop.height, 3.0) |
| cw, ch = max(1, round(crop.width * s)), max(1, round(crop.height * s)) |
| crop = crop.resize((cw, ch), Image.LANCZOS).convert("RGBA") |
| cell = Image.new("RGBA", (W, H), (236, 236, 236, 255)) |
| ox, oy = (W - cw) // 2, (H - ch) // 2 |
| cell.alpha_composite(crop, (ox, oy)) |
| dr = ImageDraw.Draw(cell) |
| rx, ry = ox + (o["x"] - x0) * s, oy + (o["y"] - y0) * s |
| dr.rectangle([rx - 2, ry - 2, rx + o["w"] * s + 1, ry + o["h"] * s + 1], outline=(230, 20, 20, 255), width=2) |
| return cell |
| |
| cell = Image.new("RGBA", (W, H), tuple(e["light"]) + (255,)) |
| big = content_area(o, decode_png(o["big"])) |
| dw, dh, _ = fit(o, (W - 16, H - 16)) |
| cell.alpha_composite(Image.fromarray(big, "RGBA").resize((dw, dh), Image.LANCZOS), ((W - dw) // 2, (H - dh) // 2)) |
| ImageDraw.Draw(cell).text((4, H - 14), "(on its light background)", fill=(90, 90, 90), font=F11) |
| return cell |
|
|
|
|
| def draw_sheet(rows, typ, shots, path): |
| g = GEOM[typ] |
| cw, ch = g["cell"] |
| xw, xh = g["ctx"] |
| pad, lab, cap = 6, 30, 16 |
| W = lab + xw + 4 * (cw + pad) + pad |
| rh = max(ch, xh) + cap + 10 |
| H = 26 + rh * len(rows) |
| sheet = Image.new("RGB", (W, H), (255, 255, 255)) |
| dr = ImageDraw.Draw(sheet) |
| heads = ["light page (red box)", "A keep", "B invert", "C adapt", "D remove"] |
| xs = [lab, lab + xw + pad] + [lab + xw + pad + k * (cw + pad) for k in range(1, 4)] |
| for x, t in zip(xs, heads): |
| dr.text((x + 4, 6), t, fill=(0, 0, 0), font=F13) |
| meta = [] |
| for r, e in enumerate(rows): |
| y = 26 + r * rh |
| o = e["o"] |
| dr.text((4, y + ch // 2 - 8), f"#{r + 1}", fill=(0, 0, 0), font=F15) |
| sheet.paste(context_cell(shots.get(e["page"]), e, (xw, xh)).convert("RGB"), (xs[0], y)) |
| big = content_area(o, decode_png(o["big"])) |
| offers = "AB" |
| inv = invert(big) |
| ad = adapt(big) |
| cells = [variant_cell(big, o, typ, (cw, ch), e["dsurf"], e["dtext"], "A"), variant_cell(inv, o, typ, (cw, ch), e["dsurf"], e["dtext"], "B")] |
| if ad is not None: |
| offers += "C" |
| cells.append(variant_cell(ad, o, typ, (cw, ch), e["dsurf"], e["dtext"], "C")) |
| else: |
| cells.append(None) |
| if typ == LARGE and o["kind"] == "bg": |
| offers += "D" |
| cells.append(variant_cell(big, o, typ, (cw, ch), e["dsurf"], e["dtext"], "D")) |
| else: |
| cells.append(None) |
| for k, c in enumerate(cells): |
| x = xs[1 + k] |
| if c is None: |
| dr.rectangle([x, y, x + cw - 1, y + ch - 1], fill=(225, 225, 225)) |
| dr.text((x + cw // 2 - 14, y + ch // 2 - 7), "n/a", fill=(140, 140, 140), font=F13) |
| else: |
| sheet.paste(c.convert("RGB"), (x, y)) |
| flags = [f for f, v in (("logo", o.get("logo")), ("links home", o.get("homeLink")), ("in a link", o.get("inLink") and not o.get("homeLink")), ("in header", o.get("header"))) if v] |
| _, _, z = fit(o, (cw - 12, ch - 12)) |
| zoom = f", shown {z:.1f}x" if abs(z - 1) > 0.15 else "" |
| dr.text((xs[1], y + ch + 2), f"#{r + 1}: {KIND.get(o['kind'], o['kind'])} {o['w']}x{o['h']} px{zoom}" + (" · " + ", ".join(flags) if flags else ""), fill=(60, 60, 60), font=F11) |
| meta.append({"k": r + 1, "page": e["page"], "crawl": e["crawl"], "img": e["img"], "bucket": e["bucket"], "offer": offers, "kind": o["kind"]}) |
| sheet.save(path, optimize=True) |
| return meta |
|
|
|
|
| def render_chunk(job): |
| wid, rows_by_sheet, crawls, out = job |
| shots = {} |
| outp = Path(out) |
| lines = [] |
| for sid, typ, rows in rows_by_sheet: |
| for e in rows: |
| if e["page"] not in shots: |
| p = Path(crawls[e["crawl"]]) / "pages" / e["page"] / "shot.jpg" |
| try: |
| shots[e["page"]] = Image.open(p).convert("RGB") |
| except Exception: |
| shots[e["page"]] = None |
| try: |
| meta = draw_sheet(rows, typ, shots, outp / "sheets" / f"{sid}.png") |
| lines.append(json.dumps({"sheet": sid, "type": typ, "rows": meta})) |
| except Exception as ex: |
| print("sheet error", sid, repr(ex)[:200], flush=True) |
| if len(shots) > 64: |
| shots.clear() |
| with open(outp / f"meta.{wid}.jsonl", "a") as f: |
| f.write("\n".join(lines) + ("\n" if lines else "")) |
| return len(lines) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--crawl", action="append", required=True, help="TAG=CRAWL_DIR=HAIKU_DIR (repeatable)") |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--max-elements", type=int, default=150000) |
| ap.add_argument("--max-per-src", type=int, default=2) |
| ap.add_argument("--workers", type=int, default=28) |
| ap.add_argument("--exclude-pages", default="", help="file of page ids already done") |
| ap.add_argument("--flat-rate", type=float, default=None, help="sampling rate for opaque flat/photo-like images (mostly 'keep')") |
| args = ap.parse_args() |
| if args.flat_rate is not None: |
| RATE["flat"] = args.flat_rate |
| out = Path(args.out) |
| (out / "sheets").mkdir(parents=True, exist_ok=True) |
| crawls, tasks = {}, [] |
| done = set(l.strip() for l in open(args.exclude_pages)) if args.exclude_pages and os.path.exists(args.exclude_pages) else set() |
| for spec in args.crawl: |
| tag, cdir, hdir = spec.split("=") |
| crawls[tag] = cdir |
| for hp in Path(hdir).glob("*.json"): |
| pid = hp.stem |
| if len(pid) == 16 and pid not in done and (Path(cdir) / "pages" / pid / "light.json.gz").exists(): |
| tasks.append((cdir, hdir, pid, tag, int(hashlib.md5(pid.encode()).hexdigest()[:8], 16))) |
| random.Random(3).shuffle(tasks) |
| print(len(tasks), "pages with designs", flush=True) |
| per_src, elems, by_bucket = {}, [], {} |
| with Pool(args.workers) as pool: |
| for res in pool.imap_unordered(candidates, tasks, chunksize=16): |
| for e in res: |
| s = e["o"]["src"] |
| if per_src.get(s, 0) >= args.max_per_src: |
| continue |
| per_src[s] = per_src.get(s, 0) + 1 |
| elems.append(e) |
| by_bucket[e["bucket"]] = by_bucket.get(e["bucket"], 0) + 1 |
| if len(elems) >= args.max_elements: |
| break |
| elems = elems[: args.max_elements] |
| print(len(elems), "elements", by_bucket, flush=True) |
| with open(out / "pages_used.txt", "a") as f: |
| f.write("\n".join(sorted(set(e["page"] for e in elems))) + "\n") |
| |
| random.Random(5).shuffle(elems) |
| groups = {SMALL: [], LARGE: []} |
| for e in elems: |
| o = e["o"] |
| typ = LARGE if (o["w"] * o["h"] >= 20000 or o["w"] > 300) else SMALL |
| groups[typ].append(e) |
| sheets = [] |
| for typ, es in groups.items(): |
| n = GEOM[typ]["rows"] |
| for i in range(0, len(es), n): |
| rows = es[i:i + n] |
| sid = hashlib.sha1(json.dumps([(r["page"], r["img"]) for r in rows]).encode()).hexdigest()[:16] |
| sheets.append((sid, typ, rows)) |
| jobs = [(w, sheets[w::args.workers], crawls, str(out)) for w in range(args.workers)] |
| with Pool(args.workers) as pool: |
| n = sum(pool.map(render_chunk, jobs)) |
| print("sheets", n, flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
