| |
| """Run a trained NadirNet over held-out pages and write render jobs (Python port of nadir-infer.js decode).""" |
| import argparse, gzip, json, math, pickle |
| from pathlib import Path |
| import numpy as np |
| import torch |
|
|
| from model import NadirNet |
| from train import collate, prep |
| from mkjobs import job_from |
|
|
| TREAT = ["keep", "invert", "adapt", "remove"] |
| RENDER_T = {"keep": "keep", "invert": "invert", "adapt": "invert", "remove": "clear"} |
|
|
|
|
| from safety import oklab, lab_to_rgb, contrast, safe_treatment, enforce_contrast, is_media_keep, hue_guard, keep_invisible, rescue_surfaces, glass_panels, opaque_surfaces, dark_washes |
|
|
|
|
| def decode(L, lab, alpha_logit, img_logits, rels): |
| toks = L["tokens"][:256] |
| out, keep = [], set() |
| for i, t in enumerate(toks): |
| ain = min(1 - 1e-3, max(1e-3, t["rgba"][3])) |
| alpha = 1 / (1 + math.exp(-(alpha_logit[i] + math.log(ain / (1 - ain))))) |
| a = 1.0 if (t["rgba"][3] >= 0.999 and alpha > 0.97) else alpha |
| out.append({"rgb": lab_to_rgb(*lab[i]), "a": a, "L": float(lab[i][0])}) |
| if is_media_keep(t): |
| keep.add(i) |
| rgb = [int(x) for x in t["rgba"][:3]] |
| out[i] = {"rgb": rgb, "a": float(t["rgba"][3]), "L": oklab(*rgb)[0]} |
| else: |
| hue_guard(t, out[i]) |
| glass_panels(toks, out, keep) |
| for i, t in enumerate(toks): |
| if t["f"][24] > 0.5 and t["role"] == 0 and out[i]["L"] > 0.35: |
| o = oklab(*out[i]["rgb"]); out[i]["rgb"] = lab_to_rgb(0.2, o[1] * 0.3, o[2] * 0.3); out[i]["L"] = 0.2 |
| opaque_surfaces(toks, out, keep) |
| dark_washes(toks, out, keep) |
| keep_invisible(toks, out, rels, keep) |
| enforce_contrast(toks, out, rels, keep) |
| rescue_surfaces(toks, out, rels, keep) |
| def hx(o): |
| h = "#%02x%02x%02x" % tuple(o["rgb"]) |
| return h + ("%02x" % round(o["a"] * 255) if o["a"] < 0.999 else "") |
| colors = [None if i in keep else hx(o) for i, o in enumerate(out)] |
| treat = {} |
| for j, im in enumerate(L["images"][:32]): |
| t = TREAT[int(np.argmax(img_logits[j]))] if j < len(img_logits) else "keep" |
| treat[str(j)] = RENDER_T[t] |
| return colors, sorted(keep), treat |
|
|
|
|
| def element_treatments(model, s, lab, ctx, dev): |
| """Element head over a page's first 32 images (render jobs only know those).""" |
| m = min(32, len(s["IF"])) |
| if not m: |
| return [] |
| f = torch.from_numpy(s["IF"][:m])[None].to(dev) |
| th = torch.from_numpy(s["thumbs"][:m].transpose(0, 3, 1, 2) / 255.0).float()[None].to(dev) |
| ul = torch.from_numpy(s["IF"][:m, 23:26] * np.array([1, 0.25, 0.25], np.float32))[None].to(dev) |
| u = np.where((s["img_under"][:m] >= 0) & (s["img_under"][:m] < lab.shape[1]), s["img_under"][:m], s["root"]) |
| ud = lab[:, torch.from_numpy(u).to(dev)] |
| t, _ = model.elements(f, th, ul, ud, ctx) |
| return t[0].float().cpu().numpy().tolist() |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--data", required=True) |
| ap.add_argument("--crawl", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--n", type=int, default=150) |
| ap.add_argument("--src", type=int, default=0) |
| ap.add_argument("--tag", default="nadir") |
| ap.add_argument("--also", default="haiku,dr", help="extra render tags to emit for the same pages") |
| ap.add_argument("--haiku", default="") |
| args = ap.parse_args() |
| ck = torch.load(args.ckpt, map_location="cpu") |
| model = NadirNet(**ck["cfg"]); model.load_state_dict(ck["model"]); model.eval() |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| model.to(dev) |
| data = pickle.load(open(args.data, "rb")) |
| val = [s for s in data["val"] if s["src"] == args.src][: args.n] |
| with open(args.out, "w") as f: |
| for s in val: |
| s = prep(s) |
| L = json.load(gzip.open(Path(args.crawl) / "pages" / s["id"] / "light.json.gz")) |
| b = collate([s], dev) |
| with torch.no_grad(): |
| lab, al, ctx = model(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], torch.zeros_like(b["src"]), b["rel"], b["mask"]) |
| il = element_treatments(model, s, lab.float(), ctx, dev) |
| n = len(s["role"]) |
| colors, keep, treat = decode(L, lab[0, :n].float().cpu().numpy().tolist(), al[0, :n].float().cpu().numpy().tolist(), |
| il, s["rel"]) |
| colors = colors + [None] * (len(L["tokens"]) - len(colors)) |
| j = job_from(L, s["id"], args.tag, colors, treat) |
| j["keep"] = sorted(set(j["keep"]) | set(keep)) |
| f.write(json.dumps(j) + "\n") |
| if "dr" in args.also: |
| f.write(json.dumps(job_from(L, s["id"], "dr", [None] * len(L["tokens"]), {}, baseline=True)) + "\n") |
| if "haiku" in args.also and args.haiku: |
| hp = Path(args.haiku) / f"{s['id']}.json" |
| if hp.exists(): |
| h = json.loads(hp.read_text()) |
| outs = [h["colors"].get(str(i)) for i in range(len(L["tokens"]))] |
| f.write(json.dumps(job_from(L, s["id"], "haiku", outs, h["images"])) + "\n") |
| print("wrote", len(val), "pages") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
