500 lines
16 KiB
Python
500 lines
16 KiB
Python
"""The registered watermark scheme `wai-mark:tent64` of extensions/synthesis.md §7.5, and the
|
||
attacks of its battery (§7.5.7), written from the extension's text.
|
||
|
||
Shared by the two Python verifiers of the synthesis corpora (synthesis_verify.py and
|
||
synthesis_mark_verify.py); it shares no code with wai-rs. Standard library only.
|
||
|
||
Frames are `bytes` of RGB8, frame after frame. Every function is integer arithmetic, so its
|
||
results are the extension's to the byte.
|
||
"""
|
||
|
||
import hashlib
|
||
|
||
SCHEME = "wai-mark:tent64"
|
||
GRID = 64
|
||
CELLS = GRID * GRID
|
||
LENGTH = 248 # the shortened code's length (§7.5.3)
|
||
PARITY = 184
|
||
T = 29 # the errors the code corrects
|
||
CELLS_PER_BIT = 16
|
||
AMPLITUDE = 12 # §7.5.4
|
||
MARGIN = (3, 2) # 2·σ·T ≥ 3·Q
|
||
MIN_SIDE = 128 # §7.5.1
|
||
MAX_FRAME_BYTES = 1 << 28
|
||
MAX_INT = (1 << 53) - 1
|
||
CLASS_MAX = 3
|
||
|
||
|
||
def sha256(b):
|
||
return hashlib.sha256(b).digest()
|
||
|
||
|
||
def stream(label, seed=None):
|
||
"""§7.5.3: 32-bit big-endian words, eight from each SHA-256(label ‖ [seed ‖] n)."""
|
||
n = 0
|
||
while True:
|
||
d = hashlib.sha256(label + (b"" if seed is None else seed.to_bytes(4, "big")) + n.to_bytes(4, "big")).digest()
|
||
for o in range(0, 32, 4):
|
||
yield int.from_bytes(d[o:o + 4], "big")
|
||
n += 1
|
||
|
||
|
||
def _layout():
|
||
r = stream(b"wai-mark:tent64 stream")
|
||
perm = list(range(CELLS))
|
||
for i in range(CELLS - 1, 0, -1):
|
||
j = next(r) % (i + 1)
|
||
perm[i], perm[j] = perm[j], perm[i]
|
||
signs = [1 if next(r) % 2 == 0 else -1 for _ in range(CELLS)]
|
||
mask = [next(r) & 1 for _ in range(LENGTH)]
|
||
return perm, signs, mask
|
||
|
||
|
||
PERM, SIGNS, MASK = _layout()
|
||
BIT_CELLS = [PERM[CELLS_PER_BIT * j: CELLS_PER_BIT * j + CELLS_PER_BIT] for j in range(LENGTH)]
|
||
|
||
# ---------------------------------------------------------------- the code, §7.5.3
|
||
|
||
EXP = [0] * 510
|
||
LOG = [0] * 256
|
||
_x = 1
|
||
for _i in range(255):
|
||
EXP[_i] = _x
|
||
LOG[_x] = _i
|
||
_x <<= 1
|
||
if _x & 0x100:
|
||
_x ^= 0x11D # x^8 + x^4 + x^3 + x^2 + 1
|
||
for _i in range(255, 510):
|
||
EXP[_i] = EXP[_i - 255]
|
||
|
||
|
||
def gmul(a, b):
|
||
return 0 if a == 0 or b == 0 else EXP[LOG[a] + LOG[b]]
|
||
|
||
|
||
def generator():
|
||
"""The least common multiple of the minimal polynomials of α^1 … α^58, as an integer whose
|
||
bit d is the coefficient of x^d."""
|
||
g, cosets = 1, []
|
||
for i in range(1, 2 * T + 1):
|
||
coset, j = [], i % 255
|
||
while j not in coset:
|
||
coset.append(j)
|
||
j = 2 * j % 255
|
||
if sorted(coset) in cosets:
|
||
continue
|
||
cosets.append(sorted(coset))
|
||
p = [1] # coefficients from x^0 up
|
||
for j in coset:
|
||
q = [0] * (len(p) + 1)
|
||
for d, c in enumerate(p):
|
||
q[d + 1] ^= c
|
||
q[d] ^= gmul(c, EXP[j])
|
||
p = q
|
||
assert all(c in (0, 1) for c in p), "a minimal polynomial is binary"
|
||
pm = sum(c << d for d, c in enumerate(p))
|
||
prod, a, b = 0, g, pm # product over GF(2)
|
||
while b:
|
||
if b & 1:
|
||
prod ^= a
|
||
a <<= 1
|
||
b >>= 1
|
||
g = prod
|
||
return g
|
||
|
||
|
||
G = generator()
|
||
|
||
|
||
def encode(payload):
|
||
"""The codeword of an 8-byte payload: bit j the coefficient of x^(247 − j)."""
|
||
m = int.from_bytes(payload, "big")
|
||
r = m << PARITY
|
||
for d in range(r.bit_length() - 1, PARITY - 1, -1):
|
||
if r >> d & 1:
|
||
r ^= G << (d - PARITY)
|
||
c = (m << PARITY) | r
|
||
return [c >> (LENGTH - 1 - j) & 1 for j in range(LENGTH)]
|
||
|
||
|
||
def _syndromes(c):
|
||
return [_eval(c, i) for i in range(1, 2 * T + 1)]
|
||
|
||
|
||
def _eval(c, i):
|
||
v = 0
|
||
e = 0
|
||
while c:
|
||
if c & 1:
|
||
v ^= EXP[i * e % 255]
|
||
c >>= 1
|
||
e += 1
|
||
return v
|
||
|
||
|
||
def decode(word):
|
||
"""Bounded-distance decoding within 29: (payload, bits corrected), or None where no codeword
|
||
is within 29 bits or the payload's first byte is above 3."""
|
||
c = 0
|
||
for b in word:
|
||
c = (c << 1) | b
|
||
s = _syndromes(c)
|
||
errors = 0
|
||
if any(s):
|
||
# The error locator: the shortest linear recurrence the syndromes satisfy, over GF(2^8).
|
||
lam, prev, L, m, bd = [1], [1], 0, 1, 1
|
||
for n in range(2 * T):
|
||
d = s[n]
|
||
for i in range(1, min(L, len(lam) - 1) + 1):
|
||
d ^= gmul(lam[i], s[n - i])
|
||
if d == 0:
|
||
m += 1
|
||
continue
|
||
coef = gmul(d, EXP[(255 - LOG[bd]) % 255])
|
||
old = lam[:]
|
||
lam = lam + [0] * max(0, len(prev) + m - len(lam))
|
||
for i, bv in enumerate(prev):
|
||
lam[i + m] ^= gmul(coef, bv)
|
||
if 2 * L <= n:
|
||
L, prev, bd, m = n + 1 - L, old, d, 1
|
||
else:
|
||
m += 1
|
||
while len(lam) > 1 and lam[-1] == 0:
|
||
lam.pop()
|
||
if len(lam) - 1 != L or L > T:
|
||
return None
|
||
roots = []
|
||
for e in range(255):
|
||
v = 0
|
||
for i, co in enumerate(lam):
|
||
if co:
|
||
v ^= gmul(co, EXP[(-e * i) % 255])
|
||
if v == 0:
|
||
roots.append(e)
|
||
if len(roots) != L or any(e >= LENGTH for e in roots):
|
||
return None
|
||
for e in roots:
|
||
c ^= 1 << e
|
||
if any(_syndromes(c)):
|
||
return None
|
||
errors = L
|
||
payload = (c >> PARITY).to_bytes(8, "big")
|
||
if payload[0] > CLASS_MAX:
|
||
return None
|
||
return payload, errors
|
||
|
||
|
||
# ---------------------------------------------------------------- frames, §7.5.1 and §7.5.2
|
||
|
||
def admits(w, h, n):
|
||
return w >= MIN_SIDE and h >= MIN_SIDE and n >= 1 and w * h * 3 <= MAX_FRAME_BYTES and w * h * 3 * n <= MAX_INT
|
||
|
||
|
||
def luma(r, g, b):
|
||
return (77 * r + 150 * g + 29 * b + 128) >> 8
|
||
|
||
|
||
def tent(u, length):
|
||
c = 2 * u + 1 - length
|
||
sg = (c > 0) - (c < 0)
|
||
return sg * (length - abs(2 * abs(c) - length))
|
||
|
||
|
||
def axis(n):
|
||
"""Per sample: (cell index, tent, cell length); and the edges."""
|
||
edges = [i * n // GRID for i in range(GRID + 1)]
|
||
cell, tents, lens = [], [], []
|
||
for i in range(GRID):
|
||
L = edges[i + 1] - edges[i]
|
||
for x in range(edges[i], edges[i + 1]):
|
||
cell.append(i)
|
||
tents.append(tent(x - edges[i], L))
|
||
lens.append(L)
|
||
return cell, tents, lens, edges
|
||
|
||
|
||
def luma_plane(frame):
|
||
return [luma(frame[p], frame[p + 1], frame[p + 2]) for p in range(0, len(frame), 3)]
|
||
|
||
|
||
def bit_statistics(y, w, h):
|
||
"""T_j over one luma plane of w × h."""
|
||
cx, fx, _, _ = axis(w)
|
||
cy, fy, _, _ = axis(h)
|
||
s = [0] * CELLS
|
||
for row in range(h):
|
||
acc = [0] * GRID
|
||
base = row * w
|
||
for col in range(w):
|
||
acc[cx[col]] += fx[col] * y[base + col]
|
||
r0, f = cy[row] * GRID, fy[row]
|
||
for i in range(GRID):
|
||
s[r0 + i] += f * acc[i]
|
||
return [sum(SIGNS[k] * s[k] for k in BIT_CELLS[j]) for j in range(LENGTH)]
|
||
|
||
|
||
# ---------------------------------------------------------------- embedding, §7.5.4
|
||
|
||
def embed(unmarked, w, h, n, payload):
|
||
"""The marked frames, and each frame's amplitudes (a string of 248 digits as base-36)."""
|
||
code = encode(payload)
|
||
sigma = [1 if code[j] ^ MASK[j] else -1 for j in range(LENGTH)]
|
||
_, fx, lx, ex = axis(w)
|
||
_, fy, ly, ey = axis(h)
|
||
fb = w * h * 3
|
||
out = bytearray(unmarked)
|
||
amps = []
|
||
for f in range(n):
|
||
frame = unmarked[f * fb:(f + 1) * fb]
|
||
chosen = []
|
||
for j in range(LENGTH):
|
||
sg = sigma[j]
|
||
pixels = [] # (offset, ε·ω, z, ν)
|
||
q = 0
|
||
for k in BIT_CELLS[j]:
|
||
i, jj = k % GRID, k // GRID
|
||
eps = SIGNS[k]
|
||
for row in range(ey[jj], ey[jj + 1]):
|
||
for col in range(ex[i], ex[i + 1]):
|
||
om = fx[col] * fy[row]
|
||
q += abs(om)
|
||
pixels.append(((row * w + col) * 3, eps * om, sg * eps * om, lx[col] * ly[row]))
|
||
|
||
def at(a):
|
||
t, vals = 0, []
|
||
for p, eo, z, nu in pixels:
|
||
d = (z > 0) - (z < 0)
|
||
d *= (2 * a * abs(z) + nu) // (2 * nu)
|
||
r = min(255, max(0, frame[p] + d))
|
||
g = min(255, max(0, frame[p + 1] + d))
|
||
b = min(255, max(0, frame[p + 2] + d))
|
||
t += eo * luma(r, g, b)
|
||
vals.append((p, r, g, b))
|
||
return t, vals
|
||
|
||
pick = None
|
||
for a in range(AMPLITUDE + 1):
|
||
t, vals = at(a)
|
||
if MARGIN[1] * sg * t >= MARGIN[0] * q:
|
||
pick = (a, vals)
|
||
break
|
||
if pick is None:
|
||
pick = (AMPLITUDE, at(AMPLITUDE)[1])
|
||
a, vals = pick
|
||
for p, r, g, b in vals:
|
||
out[f * fb + p] = r
|
||
out[f * fb + p + 1] = g
|
||
out[f * fb + p + 2] = b
|
||
chosen.append(a)
|
||
amps.append(chosen)
|
||
return bytes(out), amps
|
||
|
||
|
||
# ---------------------------------------------------------------- detection, §7.5.5
|
||
|
||
def orient(y, w, h, o):
|
||
ow, oh = (h, w) if o & 4 else (w, h)
|
||
out = [0] * (w * h)
|
||
for yy in range(oh):
|
||
y1 = oh - 1 - yy if o & 2 else yy
|
||
for xx in range(ow):
|
||
x1 = ow - 1 - xx if o & 1 else xx
|
||
sx, sy = (y1, x1) if o & 4 else (x1, y1)
|
||
out[yy * ow + xx] = y[sy * w + sx]
|
||
return out, ow, oh
|
||
|
||
|
||
def read_word(samples, w, h, n, o):
|
||
fb = w * h * 3
|
||
t = [0] * LENGTH
|
||
for f in range(n):
|
||
y, ow, oh = orient(luma_plane(samples[f * fb:(f + 1) * fb]), w, h, o)
|
||
for j, v in enumerate(bit_statistics(y, ow, oh)):
|
||
t[j] += v
|
||
return [(1 if t[j] > 0 else 0) ^ MASK[j] for j in range(LENGTH)]
|
||
|
||
|
||
def detect(samples, w, h, n):
|
||
"""(payload, orientation, errors) of the first orientation that decodes, or None."""
|
||
if not admits(w, h, n) or len(samples) != w * h * 3 * n:
|
||
return None
|
||
for o in range(8):
|
||
got = decode(read_word(samples, w, h, n, o))
|
||
if got is not None:
|
||
return got[0], o, got[1]
|
||
return None
|
||
|
||
|
||
# ---------------------------------------------------------------- the battery, §7.5.7
|
||
|
||
def rnd(n, d):
|
||
"""n / d rounded to the nearest integer, halves away from zero."""
|
||
sg = (n > 0) - (n < 0)
|
||
return sg * ((2 * abs(n) + d) // (2 * d))
|
||
|
||
|
||
def clamp(v):
|
||
return 0 if v < 0 else 255 if v > 255 else v
|
||
|
||
|
||
def int_motion(frame, q):
|
||
return bytes(clamp(128 + rnd(v - 128, q) * q) for v in frame)
|
||
|
||
|
||
def upscale2(px, w, h):
|
||
"""wai-synthesis:upscale2, §4.1."""
|
||
out = bytearray(4 * w * h * 3)
|
||
ow = 2 * w
|
||
for y in range(h):
|
||
y1 = min(y + 1, h - 1)
|
||
for x in range(w):
|
||
x1 = min(x + 1, w - 1)
|
||
for c in range(3):
|
||
a, b = px[(y * w + x) * 3 + c], px[(y * w + x1) * 3 + c]
|
||
cc, d = px[(y1 * w + x) * 3 + c], px[(y1 * w + x1) * 3 + c]
|
||
out[((2 * y) * ow + 2 * x) * 3 + c] = a
|
||
out[((2 * y) * ow + 2 * x + 1) * 3 + c] = (a + b + 1) >> 1
|
||
out[((2 * y + 1) * ow + 2 * x) * 3 + c] = (a + cc + 1) >> 1
|
||
out[((2 * y + 1) * ow + 2 * x + 1) * 3 + c] = (a + b + cc + d + 2) >> 2
|
||
return bytes(out)
|
||
|
||
|
||
def denoise3(px, w, h):
|
||
"""wai-synthesis:denoise3, §4.2."""
|
||
k = ((1, 2, 1), (2, 4, 2), (1, 2, 1))
|
||
out = bytearray(w * h * 3)
|
||
for y in range(h):
|
||
for x in range(w):
|
||
for c in range(3):
|
||
s = 8
|
||
for j in range(3):
|
||
yy = min(max(y + j - 1, 0), h - 1)
|
||
for i in range(3):
|
||
xx = min(max(x + i - 1, 0), w - 1)
|
||
s += k[j][i] * px[(yy * w + xx) * 3 + c]
|
||
out[(y * w + x) * 3 + c] = s >> 4
|
||
return bytes(out)
|
||
|
||
|
||
def _weights(n, m):
|
||
out = []
|
||
for d in range(m):
|
||
lo, hi = d * n, (d + 1) * n
|
||
row = []
|
||
for s in range(lo // m, -(-hi // m)):
|
||
o = min(hi, (s + 1) * m) - max(lo, s * m)
|
||
if o > 0:
|
||
row.append((s, o))
|
||
out.append(row)
|
||
return out
|
||
|
||
|
||
def area(px, w, h, nw, nh):
|
||
wx, wy = _weights(w, nw), _weights(h, nh)
|
||
den = w * h
|
||
out = bytearray(nw * nh * 3)
|
||
for dy, ry in enumerate(wy):
|
||
for dx, rx in enumerate(wx):
|
||
for c in range(3):
|
||
num = 0
|
||
for sy, oy in ry:
|
||
for sx, ox in rx:
|
||
num += oy * ox * px[(sy * w + sx) * 3 + c]
|
||
out[(dy * nw + dx) * 3 + c] = (2 * num + den) // (2 * den)
|
||
return bytes(out)
|
||
|
||
|
||
ROTATIONS = {"rotate-0.5": (65534, 572), "rotate-1": (65526, 1144), "rotate-2": (65496, 2287)}
|
||
|
||
|
||
def rotate(px, w, h, cs, sn):
|
||
one = 1 << 16
|
||
|
||
def p(x, y, c):
|
||
return 0 if x < 0 or y < 0 or x >= w or y >= h else px[(y * w + x) * 3 + c]
|
||
|
||
out = bytearray(w * h * 3)
|
||
for y in range(h):
|
||
for x in range(w):
|
||
u, v = 2 * x + 1 - w, 2 * y + 1 - h
|
||
X = (cs * u + sn * v + (w - 1) * one) >> 1
|
||
Y = (-sn * u + cs * v + (h - 1) * one) >> 1
|
||
xi, fx, yi, fy = X >> 16, X & 0xFFFF, Y >> 16, Y & 0xFFFF
|
||
for c in range(3):
|
||
val = (p(xi, yi, c) * (one - fx) * (one - fy) + p(xi + 1, yi, c) * fx * (one - fy)
|
||
+ p(xi, yi + 1, c) * (one - fx) * fy + p(xi + 1, yi + 1, c) * fx * fy)
|
||
out[(y * w + x) * 3 + c] = (val + (1 << 31)) >> 32
|
||
return bytes(out)
|
||
|
||
|
||
def attack(name, px, w, h, seed):
|
||
"""One picture attack of §7.5.7 on a single frame: (result, width, height)."""
|
||
if name == "none":
|
||
return px, w, h
|
||
if name.startswith("int-motion-q"):
|
||
return int_motion(px, int(name[len("int-motion-q"):])), w, h
|
||
if name == "staged-base-q8":
|
||
hw, hh = w // 2, h // 2
|
||
half = bytearray(hw * hh * 3)
|
||
for y in range(hh):
|
||
for x in range(hw):
|
||
for c in range(3):
|
||
s = sum(px[((2 * y + dy) * w + 2 * x + dx) * 3 + c] for dy in (0, 1) for dx in (0, 1))
|
||
half[(y * hw + x) * 3 + c] = (s + 2) >> 2
|
||
return upscale2(int_motion(bytes(half), 8), hw, hh), 2 * hw, 2 * hh
|
||
if name == "denoise3":
|
||
return denoise3(px, w, h), w, h
|
||
if name == "denoise3-x4":
|
||
for _ in range(4):
|
||
px = denoise3(px, w, h)
|
||
return px, w, h
|
||
if name.startswith("scale-"):
|
||
n, d = map(int, name[len("scale-"):].split("/"))
|
||
nw, nh = w * n // d, h * n // d
|
||
return area(px, w, h, nw, nh), nw, nh
|
||
if name == "upscale2":
|
||
return upscale2(px, w, h), 2 * w, 2 * h
|
||
if name == "aspect-3/4":
|
||
nw = w * 3 // 4
|
||
return area(px, w, h, nw, h), nw, h
|
||
if name == "half-then-int-motion-q16":
|
||
nw, nh = w // 2, h // 2
|
||
return int_motion(area(px, w, h, nw, nh), 16), nw, nh
|
||
if name.startswith("crop-"):
|
||
pc = int(name[len("crop-"):-1])
|
||
dx, dy = w * pc // 200, h * pc // 200
|
||
nw, nh = w - 2 * dx, h - 2 * dy
|
||
out = b"".join(px[((y * w) + dx) * 3:((y * w) + dx + nw) * 3] for y in range(dy, dy + nh))
|
||
return out, nw, nh
|
||
if name == "mirror":
|
||
return b"".join(px[(y * w + x) * 3:(y * w + x) * 3 + 3] for y in range(h) for x in range(w - 1, -1, -1)), w, h
|
||
if name == "rotate-90":
|
||
return b"".join(px[((h - 1 - xp) * w + yp) * 3:((h - 1 - xp) * w + yp) * 3 + 3] for yp in range(w) for xp in range(h)), h, w
|
||
if name in ROTATIONS:
|
||
cs, sn = ROTATIONS[name]
|
||
return rotate(px, w, h, cs, sn), w, h
|
||
if name.startswith("noise-"):
|
||
s = int(name[len("noise-"):])
|
||
r = stream(b"wai-mark:battery noise", seed)
|
||
return bytes(clamp(v + next(r) % (2 * s + 1) - s) for v in px), w, h
|
||
if name.startswith("brightness"):
|
||
b = int(name[len("brightness"):])
|
||
return bytes(clamp(v + b) for v in px), w, h
|
||
if name.startswith("contrast-"):
|
||
n, d = map(int, name[len("contrast-"):].split("/"))
|
||
return bytes(clamp(128 + rnd((v - 128) * n, d)) for v in px), w, h
|
||
if name == "local-blur":
|
||
out = bytearray(px)
|
||
x0, x1, y0, y1 = w // 4, 3 * w // 4, h // 4, 3 * h // 4
|
||
for by in range(y0, y1, 16):
|
||
for bx in range(x0, x1, 16):
|
||
ex, ey = min(bx + 16, x1), min(by + 16, y1)
|
||
cnt = (ex - bx) * (ey - by)
|
||
for c in range(3):
|
||
tot = sum(px[(y * w + x) * 3 + c] for y in range(by, ey) for x in range(bx, ex))
|
||
mval = (2 * tot + cnt) // (2 * cnt)
|
||
for y in range(by, ey):
|
||
for x in range(bx, ex):
|
||
out[(y * w + x) * 3 + c] = mval
|
||
return bytes(out), w, h
|
||
raise ValueError("no picture attack named %r" % name)
|