File size: 7,811 Bytes
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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | """NadirNet: a relation-biased set transformer over a page's colour palette and images.
Tokens: [PAGE] + one token per (colour, role) the page uses + one token per image.
Attention logits get a learned per-head bias from the page's structure graph (text-on-surface,
surface-on-surface, border-on-surface, same-element, image-on-surface, same-colour), so the model
reasons about hierarchy and contrast the way a designer would. Outputs a dark colour per token
(OKLab, parameterised relative to the input so hue is kept unless the model chooses otherwise),
an alpha, and a compact page context.
Element head (v2): decides how every image, CSS background, sprite or pseudo-element icon on the page is shown
(keep / invert / adapt / remove) and what it is (photo, logo, icon, formula...). It reuses the trunk's image
encoder and reads the page context plus the light surface under the element and the dark surface the palette
head designed for it, so it can run for any element at any time (late-loaded icons, sprites) in a tiny second
graph without re-running the transformer.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
FT, FI, FP = 49, 41, 8
N_ROLES = 7
TREAT = ["keep", "invert", "adapt", "remove"]
IMG_ROLES = ["photo", "illustration", "diagram", "logo", "icon", "text", "backdrop", "avatar", "ui"]
N_TREAT, N_IROLE, N_CTX = len(TREAT), len(IMG_ROLES), 32
K_REL = 11
def fourier_lab(lab, n=6):
x = torch.stack([lab[..., 0], lab[..., 1] * 4, lab[..., 2] * 4], -1) # (..., 3)
freqs = (2.0 ** torch.arange(n, device=lab.device, dtype=lab.dtype)) * math.pi
ang = x[..., None] * freqs # (...,3,n)
return torch.cat([torch.sin(ang), torch.cos(ang)], -1).flatten(-2) # (..., 6n)
class MLP(nn.Module):
def __init__(self, i, h, o, drop=0.0):
super().__init__()
self.net = nn.Sequential(nn.Linear(i, h), nn.GELU(), nn.Dropout(drop), nn.Linear(h, o))
def forward(self, x):
return self.net(x)
class ThumbCNN(nn.Module):
"""32x32 RGBA thumbnail -> embedding. Alpha matters: transparent logos behave differently from photos."""
def __init__(self, o=128):
super().__init__()
c = [5, 32, 64, 96, 128]
layers = []
for a, b in zip(c[:-1], c[1:]):
layers += [nn.Conv2d(a, b, 3, padding=1), nn.GELU(), nn.Conv2d(b, b, 3, padding=1, stride=2), nn.GELU()]
self.net = nn.Sequential(*layers)
self.proj = nn.Linear(128 * 2 * 2, o)
def forward(self, x): # (B, M, 4, 32, 32) in [0,1]
B, M = x.shape[:2]
x = x.reshape(B * M, 4, 32, 32)
a = x[:, 3:4]
lum = (0.2126 * x[:, 0:1] + 0.7152 * x[:, 1:2] + 0.0722 * x[:, 2:3]) * a # premultiplied luminance
x = torch.cat([x[:, :3] * a, a, lum], 1)
x = self.net(x) # (BM,128,2,2)
return self.proj(x.flatten(1)).reshape(B, M, -1)
class Block(nn.Module):
def __init__(self, d, h, ff, drop):
super().__init__()
self.h = h
self.ln1 = nn.LayerNorm(d)
self.qkv = nn.Linear(d, 3 * d)
self.proj = nn.Linear(d, d)
self.rel = nn.Linear(K_REL, h, bias=False)
nn.init.normal_(self.rel.weight, std=0.5)
self.ln2 = nn.LayerNorm(d)
self.ff = nn.Sequential(nn.Linear(d, ff * d), nn.GELU(), nn.Dropout(drop), nn.Linear(ff * d, d))
self.drop = nn.Dropout(drop)
def forward(self, x, rel, pad_bias):
B, T, D = x.shape
q, k, v = self.qkv(self.ln1(x)).reshape(B, T, 3, self.h, D // self.h).permute(2, 0, 3, 1, 4)
bias = self.rel(rel).permute(0, 3, 1, 2) + pad_bias # (B,H,T,T)
a = F.scaled_dot_product_attention(q, k, v, attn_mask=bias.to(q.dtype))
x = x + self.drop(self.proj(a.transpose(1, 2).reshape(B, T, D)))
return x + self.drop(self.ff(self.ln2(x)))
class NadirNet(nn.Module):
def __init__(self, d=384, layers=8, heads=6, ff=4, drop=0.1):
super().__init__()
self.cfg = dict(d=d, layers=layers, heads=heads, ff=ff, drop=drop)
self.tok_in = MLP(FT + N_ROLES + 36, 2 * d, d)
self.cnn = ThumbCNN(128)
self.img_in = MLP(FI + 128, 2 * d, d)
self.page_in = MLP(FP, d, d)
self.src = nn.Embedding(2, d)
self.typ = nn.Embedding(3, d)
self.blocks = nn.ModuleList([Block(d, heads, ff, drop) for _ in range(layers)])
self.ln = nn.LayerNorm(d)
self.color_head = MLP(d + 3, d, 6)
self.ctx_head = nn.Linear(d, N_CTX)
self.elem_head = MLP(d + N_CTX + 72, d, N_TREAT + N_IROLE, drop=drop)
def embed_images(self, img_f, img_thumb):
return self.img_in(torch.cat([img_f, self.cnn(img_thumb)], -1))
def elements(self, e_f, e_thumb, under_light, under_dark, ctx):
"""e_f (B,K,FI), e_thumb (B,K,4,32,32), under_light/under_dark (B,K,3) OKLab, ctx (B,N_CTX)
-> treatment logits (B,K,N_TREAT), role logits (B,K,N_IROLE)."""
K = e_f.shape[1]
h = torch.cat([self.embed_images(e_f, e_thumb), ctx[:, None].expand(-1, K, -1), fourier_lab(under_light), fourier_lab(under_dark)], -1)
z = self.elem_head(h)
return z[..., :N_TREAT], z[..., N_TREAT:]
def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
"""tok_f (B,N,FT) tok_role (B,N) long, tok_lab (B,N,3), img_f (B,M,FI), img_thumb (B,M,4,32,32) float 0-1,
page_f (B,FP), src (B,) long, rel (B,T,T,K) with T=1+N+M, mask (B,T) bool (True=real)."""
B, N = tok_role.shape
M = img_f.shape[1]
tok = self.tok_in(torch.cat([tok_f, F.one_hot(tok_role, N_ROLES).to(tok_f.dtype), fourier_lab(tok_lab)], -1)) + self.typ.weight[1]
if M > 0:
img = self.embed_images(img_f, img_thumb) + self.typ.weight[2]
else:
img = tok.new_zeros(B, 0, tok.shape[-1])
cls = self.page_in(page_f)[:, None] + self.typ.weight[0]
x = torch.cat([cls, tok, img], 1) + self.src(src)[:, None]
pad_bias = torch.zeros(mask.shape, dtype=x.dtype, device=x.device).masked_fill(~mask, -1e4)[:, None, None, :]
for blk in self.blocks:
x = blk(x, rel, pad_bias)
x = self.ln(x)
ht = x[:, 1:1 + N]
z = self.color_head(torch.cat([ht, tok_lab], -1))
L = torch.sigmoid(z[..., 0])
s = 1.6 * torch.sigmoid(z[..., 1])
r = 0.6 * torch.tanh(z[..., 2])
a_in, b_in = tok_lab[..., 1], tok_lab[..., 2]
a = s * a_in - r * b_in + 0.12 * torch.tanh(z[..., 3])
b = r * a_in + s * b_in + 0.12 * torch.tanh(z[..., 4])
lab = torch.stack([L, a, b], -1)
alpha_logit = z[..., 5]
ctx = self.ctx_head(x[:, 0])
return lab, alpha_logit, ctx
class ElemNet(nn.Module):
"""The element head as its own graph (shares NadirNet's weights): classify any set of page images."""
def __init__(self, net):
super().__init__()
self.net = net
def forward(self, e_f, e_thumb, under_light, under_dark, ctx):
return self.net.elements(e_f, e_thumb, under_light, under_dark, ctx)
# ---------- colour maths (differentiable) ----------
OKLAB_TO_LMS = torch.tensor([[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]])
LMS_TO_RGB = torch.tensor([[4.0767416621, -3.3077115913, 0.2309699292], [-1.2684380046, 2.6097574011, -0.3413193965], [-0.0041960863, -0.7034186147, 1.7076147010]])
def oklab_to_linrgb(lab):
lms = lab @ OKLAB_TO_LMS.to(lab).T
return (lms ** 3) @ LMS_TO_RGB.to(lab).T
def luminance(lab):
rgb = oklab_to_linrgb(lab).clamp(0, 1)
return rgb @ torch.tensor([0.2126, 0.7152, 0.0722]).to(lab)
def contrast_ratio(lab1, lab2):
l1, l2 = luminance(lab1), luminance(lab2)
return (torch.maximum(l1, l2) + 0.05) / (torch.minimum(l1, l2) + 0.05)
|