File size: 5,357 Bytes
bde89e0 864f427 bde89e0 40c93db bde89e0 864f427 bde89e0 864f427 40c93db bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | #!/usr/bin/env python3
"""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"} # render.py has no adapt path
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): # shown as-is: later safety passes see its real colour
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()
|