66054c9486
live-conformance/avatar-face: the generator trained on faces the tool draws, two held-out identity priors and their tracks, the training report, the seeded generator and identity, and vectors.json (forms, frames, cases). avatar_face_verify.py runs the driver and the generator's arithmetic in Python's integers, verifies records and synthesis, and mutates the corpus. tools/wai_avatar_face_train.py renders the faces, trains, quantises and measures. Co-Authored-By: Claude Opus 5.5 <[email protected]>
978 lines
44 KiB
Python
978 lines
44 KiB
Python
#!/usr/bin/env python3
|
||
"""Train the integer face generator of `extensions/avatars.md` (§7, the lowering `face-int/1`),
|
||
measure it, and write the inputs of the corpus `live-conformance/avatar-face/`.
|
||
|
||
python tools/wai_avatar_face_train.py --out <dir> [--steps 30000] [--qat-steps 8000]
|
||
|
||
Needs torch and numpy. Everything it trains on is generated here: `render` draws a face from a
|
||
seeded parameter set (head shape, skin, hair, eyes, brows, nose, mouth, background) at a pose
|
||
and an expression (head translation, scale, roll, yaw and pitch, eyes open or shut, gaze, brows,
|
||
mouth open and its width), at four times the output size, and box-filters it down. Each picture
|
||
comes with the sixteen keypoints the parameters put on it, exactly. No picture, recording or
|
||
footage from outside this repository is read. The generated set carries this repository's
|
||
licence.
|
||
|
||
It writes, under <dir>:
|
||
|
||
- `generator.wfg`: the generator (§7.2), trained in floating point, then fine-tuned with the
|
||
integer arithmetic emulated in the loop (weights on the 2^-12 grid, activations on the
|
||
2^-15 grid, the sigmoid table), and written by rounding each float32 weight times 2^12 and
|
||
each bias times 2^27 half to even: exact from the float's bits;
|
||
- `identities/<name>.wfi`: identity priors (§7.3) of faces held out of training, and
|
||
`tracks/<name>.json`: keypoint tracks (§7.4) of held-out motion, which the corpus reads;
|
||
- `report.json`: the seeds, the steps, the losses, and the measurement of §7.9: for the
|
||
held-out set, PSNR and SSIM of the integer generator against the floating-point network and
|
||
against the rendered truth, the keypoint error of a separately trained float keypoint
|
||
detector over each, and the bitrate of the keypoint tracks under §7.4's coding.
|
||
|
||
Every seed is fixed (`--seed`), but float training is not claimed to be bit-reproducible: its
|
||
sums depend on the device and the thread count of the machine that runs it. What is pinned is
|
||
the generator's bytes, by SHA-256. The integer generation over those bytes is exact (§7.6),
|
||
whatever machine trained them. `int_forward` below is the arithmetic of §7.5 to §7.7 in numpy
|
||
int64: the integers this tool reports are those, not the float network's.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import hashlib
|
||
import json
|
||
import math
|
||
import os
|
||
import struct
|
||
import time
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
# ---------------------------------------------------------------------------------------
|
||
# The arithmetic of §7.5 (the driver's warp and keypoint maps) in numpy int64.
|
||
# ---------------------------------------------------------------------------------------
|
||
|
||
K = 16
|
||
FW = 12 # weights are integers at 2^12
|
||
FA = 15 # activations are integers at 2^15
|
||
FB = FW + FA # biases are integers at 2^27
|
||
ACT_HI = 32767
|
||
|
||
|
||
def exp_neg_q16(x: np.ndarray) -> np.ndarray:
|
||
"""integer-operators §2: exp_neg_q16 of non-positive Q16 values, elementwise."""
|
||
x = x.astype(np.int64)
|
||
u = -x
|
||
n = u // 45426
|
||
g = 45426 * n
|
||
r = x + g
|
||
t = r + 88670
|
||
t2 = t * t
|
||
q = t2 >> 16
|
||
c = 23495 * q
|
||
y = (c >> 16) + 22544
|
||
out = np.where(n >= 63, 0, y >> np.minimum(n, 62))
|
||
return out.astype(np.int64)
|
||
|
||
|
||
def rdiv(n: np.ndarray, d: np.ndarray) -> np.ndarray:
|
||
"""sgn(n)·⌊(2|n| + d) / (2d)⌋ for d ≥ 1: half away from zero (integer-operators E12)."""
|
||
a = np.abs(n)
|
||
return np.sign(n) * ((2 * a + d) // (2 * d))
|
||
|
||
|
||
def warp_and_maps(ref: np.ndarray, src: np.ndarray, drv: np.ndarray, gauss: int, bg: int):
|
||
"""§7.5: the warped reference [h, w, 3] u8 and the keypoint maps [K, h, w] in [0, 256].
|
||
|
||
`ref` is the identity's picture [h, w, 3] u8, `src` its keypoints and `drv` the driving
|
||
keypoints, each [K, 2] integers in sixteenths of a pixel (x, y)."""
|
||
h, w, _ = ref.shape
|
||
ys, xs = np.mgrid[0:h, 0:w]
|
||
px = (16 * xs + 8).astype(np.int64)[None]
|
||
py = (16 * ys + 8).astype(np.int64)[None]
|
||
dx = px - drv[:, 0].astype(np.int64)[:, None, None]
|
||
dy = py - drv[:, 1].astype(np.int64)[:, None, None]
|
||
d2 = dx * dx + dy * dy
|
||
e = (d2 * gauss + 128) >> 8
|
||
wk = exp_neg_q16(-e) # [K, h, w], each in [0, 65553]
|
||
den = bg + wk.sum(axis=0)
|
||
mx = (src[:, 0] - drv[:, 0]).astype(np.int64)[:, None, None]
|
||
my = (src[:, 1] - drv[:, 1]).astype(np.int64)[:, None, None]
|
||
ddx = rdiv((wk * mx).sum(axis=0), den)
|
||
ddy = rdiv((wk * my).sum(axis=0), den)
|
||
ux = 16 * xs + ddx
|
||
uy = 16 * ys + ddy
|
||
ix = ux // 16
|
||
iy = uy // 16
|
||
fx = ux - 16 * ix
|
||
fy = uy - 16 * iy
|
||
x0 = np.clip(ix, 0, w - 1)
|
||
x1 = np.clip(ix + 1, 0, w - 1)
|
||
y0 = np.clip(iy, 0, h - 1)
|
||
y1 = np.clip(iy + 1, 0, h - 1)
|
||
r = ref.astype(np.int64)
|
||
p00, p01, p10, p11 = r[y0, x0], r[y0, x1], r[y1, x0], r[y1, x1]
|
||
fx3, fy3 = fx[..., None], fy[..., None]
|
||
v = ((16 - fx3) * (16 - fy3) * p00 + fx3 * (16 - fy3) * p01 + (16 - fx3) * fy3 * p10 + fx3 * fy3 * p11 + 128) >> 8
|
||
return v.astype(np.uint8), (wk >> 8).astype(np.int64)
|
||
|
||
|
||
def sigmoid_table() -> np.ndarray:
|
||
"""§7.6: the `lut16` table of the mask, T[v + 32768] for v in −32768..32767."""
|
||
v = np.arange(0, 32769, dtype=np.int64)
|
||
d = 65536 + np.minimum(exp_neg_q16(-16 * v), 65536)
|
||
s = ((1 << 31) + d // 2) // d # S(v) for v ≥ 0
|
||
t = np.zeros(65536, dtype=np.int64)
|
||
t[32768:] = s[:32768] # v = 0 .. 32767
|
||
t[1:32768] = 32768 - s[32767:0:-1] # v = −32767 .. −1
|
||
t[0] = 0 # v = −32768, never read
|
||
return t
|
||
|
||
|
||
# ---------------------------------------------------------------------------------------
|
||
# The generator file (§7.2) and its integer forward pass (§7.6, §7.7) in numpy int64.
|
||
# ---------------------------------------------------------------------------------------
|
||
|
||
|
||
def layer_shapes(c1: int, c2: int, c3: int, k: int = K):
|
||
"""(cin, cout, stride) of the ten convolutions, in the file's order (§7.2)."""
|
||
return [
|
||
(6 + k, c1, 1), (c1, c2, 2), (c2, c3, 2), (c3, c3, 1), (c3, c3, 1),
|
||
(c3, 4 * c2, 1), (2 * c2, c2, 1), (c2, 4 * c1, 1), (2 * c1, c1, 1), (c1, 4, 1),
|
||
]
|
||
|
||
|
||
def write_wfg(hdr: dict, layers) -> bytes:
|
||
o = bytearray(b"WFG1")
|
||
o += struct.pack("<BBHHHHHHI", 1, hdr["k"], hdr["width"], hdr["height"], hdr["c1"], hdr["c2"], hdr["c3"], hdr["gauss"], hdr["bg"])
|
||
for w, b in layers:
|
||
o += np.asarray(w, dtype="<i2").tobytes()
|
||
o += np.asarray(b, dtype="<i4").tobytes()
|
||
return bytes(o)
|
||
|
||
|
||
def read_wfg(b: bytes):
|
||
assert b[:4] == b"WFG1" and b[4] == 1
|
||
_, k, width, height, c1, c2, c3, gauss, bg = struct.unpack_from("<BBHHHHHHI", b, 4)
|
||
hdr = dict(k=k, width=width, height=height, c1=c1, c2=c2, c3=c3, gauss=gauss, bg=bg)
|
||
at, layers = 22, []
|
||
for cin, cout, s in layer_shapes(c1, c2, c3, k):
|
||
nw = cout * cin * 9
|
||
w = np.frombuffer(b, dtype="<i2", count=nw, offset=at).astype(np.int64).reshape(cout, cin, 3, 3)
|
||
at += 2 * nw
|
||
bb = np.frombuffer(b, dtype="<i4", count=cout, offset=at).astype(np.int64)
|
||
at += 4 * cout
|
||
layers.append((w, bb, s))
|
||
assert at == len(b)
|
||
return hdr, layers
|
||
|
||
|
||
def half_up(p: np.ndarray, s: int) -> np.ndarray:
|
||
return p if s == 0 else (p + (1 << (s - 1))) >> s
|
||
|
||
|
||
def conv_int(x: np.ndarray, w: np.ndarray, b: np.ndarray, stride: int) -> np.ndarray:
|
||
"""conv 3×3, zero padding 1, then rq: pre bias, shift 12, half up (§7.6)."""
|
||
c, h, wd = x.shape
|
||
xp = np.zeros((c, h + 2, wd + 2), np.int64)
|
||
xp[:, 1:-1, 1:-1] = x
|
||
ho, wo = (h - 1) // stride + 1, (wd - 1) // stride + 1
|
||
acc = np.zeros((w.shape[0], ho, wo), np.int64)
|
||
for ky in range(3):
|
||
for kx in range(3):
|
||
patch = xp[:, ky:ky + stride * (ho - 1) + 1:stride, kx:kx + stride * (wo - 1) + 1:stride]
|
||
acc += np.einsum("oc,chw->ohw", w[:, :, ky, kx], patch)
|
||
return half_up(acc + b[:, None, None], FW)
|
||
|
||
|
||
def shuffle2(x: np.ndarray) -> np.ndarray:
|
||
c4, h, w = x.shape
|
||
c = c4 // 4
|
||
return x.reshape(c, 2, 2, h, w).transpose(0, 3, 1, 4, 2).reshape(c, 2 * h, 2 * w)
|
||
|
||
|
||
TABLE = sigmoid_table()
|
||
|
||
|
||
def int_forward(layers, warped: np.ndarray, ref: np.ndarray, heat: np.ndarray) -> np.ndarray:
|
||
"""§7.6 and §7.7 over one frame: the generated picture [h, w, 3] u8."""
|
||
relu = lambda v: np.clip(v, 0, ACT_HI) # noqa: E731
|
||
t3 = (warped.transpose(2, 0, 1).astype(np.int64) - 128) * 256
|
||
t4 = (ref.transpose(2, 0, 1).astype(np.int64) - 128) * 256
|
||
t5 = heat.astype(np.int64) * 128
|
||
x = np.concatenate([t3, t4, t5])
|
||
(w1, b1, s1), (w2, b2, s2), (w3, b3, s3), (w4, b4, _), (w5, b5, _), (w6, b6, _), (w7, b7, _), (w8, b8, _), (w9, b9, _), (w10, b10, _) = layers
|
||
e1 = relu(conv_int(x, w1, b1, s1))
|
||
e2 = relu(conv_int(e1, w2, b2, s2))
|
||
e3 = relu(conv_int(e2, w3, b3, s3))
|
||
m1 = relu(conv_int(e3, w4, b4, 1))
|
||
r = relu(conv_int(m1, w5, b5, 1) + e3)
|
||
u2 = relu(shuffle2(conv_int(r, w6, b6, 1)))
|
||
d2 = relu(conv_int(np.concatenate([u2, e2]), w7, b7, 1))
|
||
u1 = relu(shuffle2(conv_int(d2, w8, b8, 1)))
|
||
d1 = relu(conv_int(np.concatenate([u1, e1]), w9, b9, 1))
|
||
o = conv_int(d1, w10, b10, 1)
|
||
g = np.clip(o[:3], -32768, 32767)
|
||
z = np.clip(o[3:4], -32767, 32767)
|
||
m = TABLE[z + 32768]
|
||
a = half_up(t3 * m, 15)
|
||
om = 32768 - m
|
||
bb = half_up(g * om, 15)
|
||
s = a + bb
|
||
out = np.clip(half_up(s, 8) + 128, 0, 255)
|
||
return out.transpose(1, 2, 0).astype(np.uint8)
|
||
|
||
|
||
# ---------------------------------------------------------------------------------------
|
||
# The keypoint track coding (§7.4): what the bitrate measurement counts.
|
||
# ---------------------------------------------------------------------------------------
|
||
|
||
|
||
def zigzag(v: int) -> int:
|
||
return 2 * v if v >= 0 else -2 * v - 1
|
||
|
||
|
||
def eg_bits(u: int, k: int) -> int:
|
||
"""The length of Exp-Golomb order k of u."""
|
||
m = (u >> k) + 1
|
||
return 2 * m.bit_length() - 1 + k
|
||
|
||
|
||
def encode_track(frames: np.ndarray, step: int, gen: bytes, ident: bytes, first: int, rate=(25, 1)) -> bytes:
|
||
"""A `WFK1` payload of the frames [n, K, 2] (sixteenths of a pixel), the coordinates
|
||
rounded to multiples of `step` sixteenths, with the order of the deltas' code chosen to
|
||
make the payload shortest."""
|
||
q = np.floor(frames / step + 0.5).astype(np.int64)
|
||
n = q.shape[0]
|
||
deltas = [zigzag(int(v)) for v in (q[1:] - q[:-1]).reshape(-1)]
|
||
best = min(range(16), key=lambda k: sum(eg_bits(u, k) for u in deltas))
|
||
out = bytearray(b"WFK1")
|
||
out += struct.pack("<BBBB", 1, q.shape[1], step, best)
|
||
out += gen + ident
|
||
out += struct.pack("<IHHH", first, n, rate[0], rate[1])
|
||
out += q[0].reshape(-1).astype("<i2").tobytes()
|
||
bits = []
|
||
for u in deltas:
|
||
m = (u >> best) + 1
|
||
nb = m.bit_length()
|
||
bits += [0] * (nb - 1)
|
||
bits += [int(c) for c in format(m, "b")]
|
||
if best:
|
||
bits += [int(c) for c in format(u & ((1 << best) - 1), "0%db" % best)]
|
||
while len(bits) % 8:
|
||
bits.append(0)
|
||
for i in range(0, len(bits), 8):
|
||
out.append(int("".join(map(str, bits[i:i + 8])), 2))
|
||
return bytes(out)
|
||
|
||
|
||
# ---------------------------------------------------------------------------------------
|
||
# The generated faces.
|
||
# ---------------------------------------------------------------------------------------
|
||
|
||
KEYPOINTS = [
|
||
"left eye", "right eye", "left upper lid", "right upper lid", "left pupil", "right pupil",
|
||
"nose tip", "mouth left", "mouth right", "upper lip", "lower lip", "chin",
|
||
"left brow", "right brow", "left cheek", "right cheek",
|
||
]
|
||
|
||
|
||
def identity_params(rng: np.random.Generator) -> dict:
|
||
skin = np.array([rng.uniform(0.35, 0.95), 0, 0])
|
||
skin[1] = skin[0] * rng.uniform(0.62, 0.82)
|
||
skin[2] = skin[0] * rng.uniform(0.45, 0.70)
|
||
hair = rng.uniform(0.02, 0.6) * np.array([1.0, rng.uniform(0.6, 0.95), rng.uniform(0.3, 0.8)])
|
||
if rng.random() < 0.12:
|
||
hair = rng.uniform(0.6, 0.85) * np.ones(3)
|
||
return dict(
|
||
bg0=rng.uniform(0.1, 0.9, 3), bg1=rng.uniform(0.1, 0.9, 3), bgdir=rng.uniform(0, 2 * math.pi),
|
||
bgfreq=rng.uniform(2, 9), bgamp=rng.uniform(0, 0.08),
|
||
skin=skin, hair=hair,
|
||
a=rng.uniform(0.25, 0.31), b=rng.uniform(0.32, 0.38),
|
||
hairline=rng.uniform(-0.62, -0.38), hairvol=rng.uniform(0.0, 0.10), sidehair=rng.uniform(0, 1),
|
||
eye_x=rng.uniform(0.34, 0.44), eye_y=rng.uniform(0.0, 0.14), eye_w=rng.uniform(0.40, 0.52),
|
||
eye_h=rng.uniform(0.020, 0.030), iris=rng.uniform(0.1, 0.6, 3) * np.array([1, 1, 1.3]),
|
||
iris_r=rng.uniform(0.85, 1.15), brow_t=rng.uniform(0.008, 0.016), brow_y=rng.uniform(0.16, 0.24),
|
||
nose_l=rng.uniform(0.20, 0.30), nose_w=rng.uniform(0.05, 0.08),
|
||
mouth_y=rng.uniform(0.46, 0.58), mouth_w=rng.uniform(0.55, 0.78), lip=np.clip(skin * np.array([rng.uniform(0.75, 1.0), 0.55, 0.55]), 0, 1),
|
||
glasses=rng.random() < 0.2, gl_col=rng.uniform(0.0, 0.5, 3),
|
||
neck=rng.uniform(0.35, 0.55), shirt=rng.uniform(0.05, 0.95, 3),
|
||
)
|
||
|
||
|
||
def pose_params(rng: np.random.Generator, neutral: bool = False) -> dict:
|
||
s = 0.35 if neutral else 1.0
|
||
return dict(
|
||
tx=rng.uniform(-0.06, 0.06) * s, ty=rng.uniform(-0.05, 0.05) * s, scale=1 + rng.uniform(-0.08, 0.08) * s,
|
||
roll=rng.uniform(-0.25, 0.25) * s, yaw=rng.uniform(-0.45, 0.45) * s, pitch=rng.uniform(-0.25, 0.25) * s,
|
||
open=0.0 if neutral and rng.random() < 0.7 else float(np.clip(rng.uniform(-0.3, 1.0), 0, 1)),
|
||
smile=rng.uniform(-0.2, 0.45) * s,
|
||
blink_l=0.0 if neutral else float(np.clip(rng.uniform(-1.5, 1.0), 0, 1)),
|
||
blink_r=None if rng.random() < 0.85 else float(np.clip(rng.uniform(-1.0, 1.0), 0, 1)),
|
||
gx=rng.uniform(-1, 1) * s, gy=rng.uniform(-0.7, 0.7) * s, brow=rng.uniform(-1, 1) * s,
|
||
)
|
||
|
||
|
||
def rot(roll, yaw, pitch):
|
||
cr, sr, cy, sy, cp, sp = math.cos(roll), math.sin(roll), math.cos(yaw), math.sin(yaw), math.cos(pitch), math.sin(pitch)
|
||
rz = np.array([[cr, -sr, 0], [sr, cr, 0], [0, 0, 1]])
|
||
ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
|
||
rx = np.array([[1, 0, 0], [0, cp, -sp], [0, sp, cp]])
|
||
return rz @ ry @ rx
|
||
|
||
|
||
def render(idp: dict, pp: dict, size: int = 64, ss: int = 4):
|
||
"""The face of `idp` at the pose and expression `pp`: [size, size, 3] u8 and its
|
||
keypoints [K, 2] in pixels (x, y) of the output."""
|
||
n = size * ss
|
||
yy, xx = np.mgrid[0:n, 0:n].astype(np.float64)
|
||
u = (xx + 0.5) / n - 0.5
|
||
v = (yy + 0.5) / n - 0.5
|
||
img = np.zeros((n, n, 3))
|
||
# Background: a gradient and a soft grating.
|
||
d = math.cos(idp["bgdir"]) * u + math.sin(idp["bgdir"]) * v
|
||
tt = np.clip(d + 0.5, 0, 1)[..., None]
|
||
img[:] = idp["bg0"] * (1 - tt) + idp["bg1"] * tt
|
||
img += idp["bgamp"] * np.sin(2 * math.pi * idp["bgfreq"] * (u + 0.7 * v))[..., None]
|
||
|
||
R = rot(pp["roll"], pp["yaw"], pp["pitch"])
|
||
sc = pp["scale"]
|
||
cx, cy = pp["tx"], pp["ty"] - 0.02
|
||
A, B, C = idp["a"], idp["b"], idp["a"] * 0.95
|
||
|
||
def proj(p):
|
||
q = R @ np.asarray(p, dtype=np.float64)
|
||
return np.array([cx + sc * q[0], cy + sc * q[1]]), q[2]
|
||
|
||
def surf(x, y):
|
||
z2 = 1 - (x / A) ** 2 - (y / B) ** 2
|
||
return C * math.sqrt(max(z2, 0.0))
|
||
|
||
def ell(c0, ax, ay, ang=0.0):
|
||
ca, sa = math.cos(ang), math.sin(ang)
|
||
du, dv = u - c0[0], v - c0[1]
|
||
ru = (ca * du + sa * dv) / max(ax, 1e-6)
|
||
rv = (-sa * du + ca * dv) / max(ay, 1e-6)
|
||
return ru * ru + rv * rv
|
||
|
||
def paint(mask, col, soft=1.5 / n):
|
||
nonlocal img
|
||
m = np.clip(mask, 0, 1)[..., None]
|
||
img = img * (1 - m) + np.asarray(col) * m
|
||
|
||
def inside(q, edge):
|
||
# q < 1 is inside; a soft edge of about `edge` in the normalised measure.
|
||
return np.clip((1 - q) / edge + 0.5, 0, 1)
|
||
|
||
roll = pp["roll"]
|
||
# Neck and shirt.
|
||
nk, _ = proj([0, B * 0.85, -C * 0.2])
|
||
paint(inside(ell(nk + np.array([0, 0.12]), idp["neck"] * A * sc, 0.25 * sc, roll * 0.3), 0.08), idp["skin"] * 0.82)
|
||
paint(np.clip((v - (nk[1] + 0.22 * sc)) / 0.01, 0, 1) * inside(ell([nk[0], 0.62], 0.5, 0.3), 0.1), idp["shirt"])
|
||
# Hair volume behind the head.
|
||
hc, _ = proj([0, -B * 0.25, -C * 0.3])
|
||
paint(inside(ell(hc, (A + idp["hairvol"] + 0.02 * idp["sidehair"]) * sc, (B * 0.95 + idp["hairvol"]) * sc, roll), 0.06), idp["hair"])
|
||
# The head: the silhouette of the rotated ellipsoid, approximated by its projected centre
|
||
# and axes, and shaded by a light the head turns under.
|
||
hcen, _ = proj([0, 0, 0])
|
||
q = ell(hcen, A * sc, B * sc, roll)
|
||
head = inside(q, 0.05)
|
||
du, dv = u - hcen[0], v - hcen[1]
|
||
ca, sa = math.cos(roll), math.sin(roll)
|
||
lu = (ca * du + sa * dv) / (A * sc)
|
||
lv = (-sa * du + ca * dv) / (B * sc)
|
||
nz = np.sqrt(np.clip(1 - lu * lu - lv * lv, 0, 1))
|
||
light = np.array([-0.45, -0.5, 0.74])
|
||
lam = np.clip(lu * light[0] * math.cos(pp["yaw"]) + lv * light[1] + nz * light[2] + 0.25 * math.sin(pp["yaw"]), 0, 1)
|
||
face_col = idp["skin"][None, None] * (0.55 + 0.5 * lam[..., None])
|
||
m = head[..., None]
|
||
img = img * (1 - m) + face_col * m
|
||
# Hair over the top of the head, above a hairline that turns with it.
|
||
hl = idp["hairline"] + 0.08 * pp["pitch"]
|
||
hairmask = head * np.clip((hl - lv) / 0.06 + 0.5, 0, 1)
|
||
side = idp["sidehair"] * np.clip((np.abs(lu) - 0.82) / 0.05, 0, 1) * np.clip((0.2 - lv) / 0.1, 0, 1)
|
||
paint(np.maximum(hairmask, head * side), idp["hair"] * (0.8 + 0.3 * lam[..., None]))
|
||
|
||
kp = np.zeros((K, 2))
|
||
|
||
def visible_w(xside):
|
||
# Foreshortening of a feature at x on the surface: |n·view| for a turned head.
|
||
return max(0.25, math.cos(pp["yaw"] + math.asin(max(-0.95, min(0.95, xside / A)))))
|
||
|
||
# Eyes.
|
||
ey = -idp["eye_y"] * B
|
||
for side, (ki, li, pi_) in ((-1, (0, 2, 4)), (1, (1, 3, 5))):
|
||
ex = side * idp["eye_x"] * A
|
||
p3 = [ex, ey, surf(ex, ey) - 0.01]
|
||
c2, _ = proj(p3)
|
||
fw = visible_w(ex)
|
||
blink = pp["blink_l"] if (side < 0 or pp["blink_r"] is None) else pp["blink_r"]
|
||
ew = idp["eye_w"] * A * sc * fw
|
||
eh = idp["eye_h"] * sc * (1 - 0.92 * blink)
|
||
socket = ell(c2, ew * 0.75, idp["eye_h"] * sc * 1.1, roll)
|
||
paint(inside(socket, 0.3) * 0.25, idp["skin"] * 0.6)
|
||
eyem = inside(ell(c2, ew * 0.5, eh, roll), 0.12)
|
||
paint(eyem, [0.95, 0.94, 0.92])
|
||
g = np.array([pp["gx"] * ew * 0.22, pp["gy"] * eh * 0.35])
|
||
g = np.array([ca * g[0] - sa * g[1], sa * g[0] + ca * g[1]])
|
||
ic = c2 + g
|
||
ir = idp["iris_r"] * idp["eye_h"] * sc
|
||
paint(inside(ell(ic, ir * fw, ir, roll), 0.15) * eyem, idp["iris"])
|
||
paint(inside(ell(ic, ir * 0.45 * fw, ir * 0.45, roll), 0.2) * eyem, [0.03, 0.03, 0.04])
|
||
lid = c2 + np.array([sa * eh, -ca * eh])
|
||
paint(inside(ell(c2, ew * 0.52, eh * 1.06, roll), 0.12) * (1 - inside(ell(c2, ew * 0.48, eh * 0.9, roll), 0.12)), idp["skin"] * 0.35)
|
||
kp[ki] = c2
|
||
kp[li] = lid
|
||
kp[pi_] = ic
|
||
# Brows.
|
||
bx = ex
|
||
by = ey - idp["brow_y"] * B - 0.012 * pp["brow"]
|
||
b2, _ = proj([bx, by, surf(bx, by)])
|
||
paint(inside(ell(b2, ew * 0.62, idp["brow_t"] * sc, roll + side * 0.08 * (pp["brow"] - 0.3)), 0.25), idp["hair"] * 0.7)
|
||
kp[12 if side < 0 else 13] = b2
|
||
# Glasses.
|
||
if idp["glasses"]:
|
||
for ki in (0, 1):
|
||
r0 = idp["eye_w"] * A * sc * 0.62
|
||
ring = ell(kp[ki], r0, r0 * 0.8, roll)
|
||
paint(np.clip(1 - np.abs(ring - 1) / 0.12, 0, 1), idp["gl_col"])
|
||
# Nose: shading along its ridge and two nostrils, with the tip standing out of the face.
|
||
ny = idp["nose_l"] * B
|
||
tip, _ = proj([0, ny, surf(0, ny) + 0.05])
|
||
root, _ = proj([0, ey, surf(0, ey)])
|
||
paint(inside(ell((tip + root) / 2, idp["nose_w"] * 0.35 * sc, np.linalg.norm(tip - root) / 2 + 0.01, roll), 0.4) * 0.25, idp["skin"] * 0.7)
|
||
for side in (-1, 1):
|
||
nn, _ = proj([side * idp["nose_w"] * 0.45, ny + 0.008, surf(0, ny) + 0.03])
|
||
paint(inside(ell(nn, 0.010 * sc, 0.006 * sc, roll), 0.4), idp["skin"] * 0.3)
|
||
kp[6] = tip
|
||
# Mouth: lips, the opening and teeth.
|
||
my = idp["mouth_y"] * B
|
||
mw = idp["mouth_w"] * A * (1 + 0.35 * pp["smile"])
|
||
lc, _ = proj([-mw / 2, my - 0.012 * pp["smile"], surf(-mw / 2, my)])
|
||
rc, _ = proj([mw / 2, my - 0.012 * pp["smile"], surf(mw / 2, my)])
|
||
mc, _ = proj([0, my, surf(0, my) + 0.01])
|
||
opn = pp["open"] * 0.045
|
||
up, _ = proj([0, my - 0.008 - opn * 0.2, surf(0, my) + 0.012])
|
||
lo, _ = proj([0, my + 0.009 + opn, surf(0, my) + 0.01])
|
||
mid = (up + lo) / 2
|
||
half_w = np.linalg.norm(rc - lc) / 2
|
||
ang = math.atan2(rc[1] - lc[1], rc[0] - lc[0])
|
||
hh = np.linalg.norm(lo - up) / 2
|
||
paint(inside(ell(mid, half_w, hh + 0.005, ang), 0.15), idp["lip"])
|
||
if pp["open"] > 0.05:
|
||
inner = inside(ell(mid, half_w * 0.85, max(hh - 0.006, 0.002), ang), 0.2)
|
||
paint(inner, [0.25, 0.06, 0.07])
|
||
teeth = inner * np.clip(((up[1] + 0.005 * sc + 0.01 * pp["open"]) - v) / 0.004 + 0.5, 0, 1)
|
||
paint(teeth, [0.93, 0.92, 0.88])
|
||
kp[7], kp[8], kp[9], kp[10] = lc, rc, up, lo
|
||
# Chin and cheeks.
|
||
chin, _ = proj([0, B * 0.98, surf(0, B * 0.98) + 0.01])
|
||
kp[11] = chin
|
||
for ki, side in ((14, -1), (15, 1)):
|
||
cpt, _ = proj([side * A * 0.97, 0.15 * B, surf(side * A * 0.97, 0.15 * B)])
|
||
kp[ki] = cpt
|
||
img = np.clip(img, 0, 1)
|
||
small = img.reshape(size, ss, size, ss, 3).mean(axis=(1, 3))
|
||
out = np.clip(np.floor(small * 255 + 0.5), 0, 255).astype(np.uint8)
|
||
kpx = (kp + 0.5) * size # u, v in [-0.5, 0.5) -> pixels
|
||
return out, kpx
|
||
|
||
|
||
def kp_q4(kp: np.ndarray) -> np.ndarray:
|
||
"""Keypoints in pixels to sixteenths of a pixel, rounded half to even."""
|
||
return np.round(kp * 16).astype(np.int64)
|
||
|
||
|
||
# ---------------------------------------------------------------------------------------
|
||
# Training (torch).
|
||
# ---------------------------------------------------------------------------------------
|
||
|
||
|
||
def make_sample(seed: int, size: int, gauss: int, bg: int):
|
||
rng = np.random.default_rng(seed)
|
||
idp = identity_params(rng)
|
||
ref, rk = render(idp, pose_params(rng, neutral=True), size)
|
||
tgt, tk = render(idp, pose_params(rng), size)
|
||
rq, tq = kp_q4(rk), kp_q4(tk)
|
||
warped, heat = warp_and_maps(ref, rq, tq, gauss, bg)
|
||
return ref, warped, heat, tgt, rq, tq
|
||
|
||
|
||
def sequence(idp_seed: int, n: int, size: int, fps: int = 25):
|
||
"""A held-out motion: smooth head motion, blinks, speech-like mouth, glances."""
|
||
rng = np.random.default_rng(idp_seed)
|
||
idp = identity_params(rng)
|
||
t = np.arange(n) / fps
|
||
|
||
def smooth(amp, fmin=0.15, fmax=0.6):
|
||
f = rng.uniform(fmin, fmax, 3)
|
||
ph = rng.uniform(0, 2 * math.pi, 3)
|
||
return amp * sum(np.sin(2 * math.pi * f[i] * t + ph[i]) for i in range(3)) / 3
|
||
|
||
blink = np.zeros(n)
|
||
for b0 in rng.uniform(0, n / fps, max(1, int(n / fps / 3))):
|
||
blink = np.maximum(blink, np.clip(1 - np.abs(t - b0) / 0.12, 0, 1))
|
||
speech = np.clip(0.5 + 0.5 * np.sin(2 * math.pi * rng.uniform(3, 5) * t + rng.uniform(0, 6)) * (0.6 + 0.4 * np.sin(2 * math.pi * 0.7 * t)), 0, 1)
|
||
poses = []
|
||
tx, ty, sc, ro, yw, pi_ = smooth(0.05), smooth(0.04), 1 + smooth(0.06), smooth(0.18), smooth(0.35), smooth(0.18)
|
||
sm, gx, gy, br = smooth(0.25) + 0.1, smooth(0.8, 0.2, 1.0), smooth(0.5, 0.2, 1.0), smooth(0.7, 0.2, 0.8)
|
||
for i in range(n):
|
||
poses.append(dict(tx=tx[i], ty=ty[i], scale=sc[i], roll=ro[i], yaw=yw[i], pitch=pi_[i], open=float(speech[i]) * 0.9,
|
||
smile=sm[i], blink_l=float(blink[i]), blink_r=None, gx=gx[i], gy=gy[i], brow=br[i]))
|
||
rng2 = np.random.default_rng(idp_seed + 1)
|
||
ref, rk = render(idp, pose_params(rng2, neutral=True), size)
|
||
frames, kps = zip(*(render(idp, p, size) for p in poses))
|
||
return ref, rk, np.stack(frames), np.stack(kps)
|
||
|
||
|
||
try:
|
||
import torch as _torch
|
||
from torch.utils.data import IterableDataset as _Iterable
|
||
except ImportError: # the measurement-free parts of this tool run without torch
|
||
_torch, _Iterable = None, object
|
||
|
||
|
||
POSES = 9 # renders per training identity: a near-neutral one and eight others
|
||
|
||
|
||
def _render_identity(args):
|
||
seed, size = args
|
||
rng = np.random.default_rng(seed)
|
||
idp = identity_params(rng)
|
||
imgs, kps = [], []
|
||
for j in range(POSES):
|
||
img, kp = render(idp, pose_params(rng, neutral=(j == 0)), size)
|
||
imgs.append(img)
|
||
kps.append(kp)
|
||
return np.stack(imgs), np.stack(kps)
|
||
|
||
|
||
def render_pool(path: Path, n_id: int, size: int, workers: int, base: int):
|
||
"""Render `n_id` training identities, POSES each, into `imgs.npy` and `kps.npy` (memory
|
||
mapped, so the data loader's workers share them)."""
|
||
import multiprocessing as mp
|
||
|
||
imgs = np.lib.format.open_memmap(path / "imgs.npy", mode="w+", dtype=np.uint8, shape=(n_id, POSES, size, size, 3))
|
||
kps = np.lib.format.open_memmap(path / "kps.npy", mode="w+", dtype=np.float64, shape=(n_id, POSES, K, 2))
|
||
t0 = time.time()
|
||
with mp.get_context("spawn").Pool(workers) as pool:
|
||
for i, (im, kp) in enumerate(pool.imap(_render_identity, [(base + i, size) for i in range(n_id)], chunksize=16)):
|
||
imgs[i] = im
|
||
kps[i] = kp
|
||
if (i + 1) % 2000 == 0:
|
||
print(f"rendered {i + 1}/{n_id} identities ({time.time() - t0:.0f} s)", flush=True)
|
||
imgs.flush()
|
||
kps.flush()
|
||
del imgs, kps
|
||
|
||
|
||
class FaceStream(_Iterable):
|
||
"""Training samples, made in the data loader's workers from the rendered pool: a
|
||
reference face, its warp to another pose of the same face, the keypoint maps, and the
|
||
face rendered at that pose."""
|
||
|
||
def __init__(self, pool: Path, base, gauss, bg):
|
||
self.pool, self.base, self.gauss, self.bg = str(pool), base, gauss, bg
|
||
|
||
def __iter__(self):
|
||
info = _torch.utils.data.get_worker_info()
|
||
wid = info.id if info else 0
|
||
imgs = np.load(self.pool + "/imgs.npy", mmap_mode="r")
|
||
kps = np.load(self.pool + "/kps.npy", mmap_mode="r")
|
||
rng = np.random.default_rng(self.base + 1_000_003 * wid)
|
||
f = lambda a: _torch.from_numpy(np.ascontiguousarray(a.transpose(2, 0, 1)).astype(np.float32)) # noqa: E731
|
||
while True:
|
||
i = int(rng.integers(0, imgs.shape[0]))
|
||
r = 0 if rng.random() < 0.6 else int(rng.integers(1, POSES))
|
||
t = int(rng.integers(0, POSES - 1))
|
||
t = t + 1 if t >= r else t
|
||
ref, tgt = np.asarray(imgs[i, r]), np.asarray(imgs[i, t])
|
||
warped, heat = warp_and_maps(ref, kp_q4(np.asarray(kps[i, r])), kp_q4(np.asarray(kps[i, t])), self.gauss, self.bg)
|
||
yield f(warped), f(ref), _torch.from_numpy(heat.astype(np.float32)), f(tgt)
|
||
|
||
|
||
def train(args, out: Path, report: dict):
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
dev = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
|
||
torch.manual_seed(args.seed)
|
||
size, gauss, bg = args.size, args.gauss, args.bg
|
||
c1, c2, c3 = args.c1, args.c2, args.c3
|
||
shapes = layer_shapes(c1, c2, c3)
|
||
|
||
class Q(torch.autograd.Function):
|
||
"""Round half up onto a grid, gradient straight through."""
|
||
|
||
@staticmethod
|
||
def forward(ctx, x, scale):
|
||
return torch.floor(x * scale + 0.5) / scale
|
||
|
||
@staticmethod
|
||
def backward(ctx, g):
|
||
return g, None
|
||
|
||
class QE(torch.autograd.Function):
|
||
"""Round half to even onto a grid (the export's rule), gradient straight through."""
|
||
|
||
@staticmethod
|
||
def forward(ctx, x, scale):
|
||
return torch.round(x * scale) / scale
|
||
|
||
@staticmethod
|
||
def backward(ctx, g):
|
||
return g, None
|
||
|
||
class Net(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.convs = nn.ModuleList([nn.Conv2d(ci, co, 3, stride=s, padding=1) for ci, co, s in shapes])
|
||
self.qat = False
|
||
|
||
def conv(self, i, x):
|
||
c = self.convs[i]
|
||
w, b = c.weight, c.bias
|
||
if self.qat:
|
||
w = QE.apply(torch.clamp(w, -32768 / 4096, 32767 / 4096), 4096.0)
|
||
b = QE.apply(b, float(1 << FB))
|
||
y = F.conv2d(x, w, b, stride=c.stride, padding=1)
|
||
if self.qat:
|
||
y = Q.apply(y, 32768.0)
|
||
return y
|
||
|
||
def forward(self, warped, ref, heat):
|
||
act = lambda v: torch.clamp(v, 0, ACT_HI / 32768) # noqa: E731
|
||
t3 = (warped - 128) / 128
|
||
t4 = (ref - 128) / 128
|
||
t5 = heat / 256
|
||
x = torch.cat([t3, t4, t5], 1)
|
||
e1 = act(self.conv(0, x))
|
||
e2 = act(self.conv(1, e1))
|
||
e3 = act(self.conv(2, e2))
|
||
m1 = act(self.conv(3, e3))
|
||
r = act(self.conv(4, m1) + e3)
|
||
u2 = act(F.pixel_shuffle(self.conv(5, r), 2))
|
||
d2 = act(self.conv(6, torch.cat([u2, e2], 1)))
|
||
u1 = act(F.pixel_shuffle(self.conv(7, d2), 2))
|
||
d1 = act(self.conv(8, torch.cat([u1, e1], 1)))
|
||
o = self.conv(9, d1)
|
||
g = torch.clamp(o[:, :3], -1, 32767 / 32768)
|
||
z = o[:, 3:4]
|
||
m = torch.sigmoid(8 * z)
|
||
if self.qat:
|
||
m = Q.apply(m, 32768.0)
|
||
s = m * t3 + (1 - m) * g
|
||
return s * 128 + 128, g * 128 + 128, 8 * z
|
||
|
||
pool = out / "pool"
|
||
pool.mkdir(exist_ok=True)
|
||
if not (pool / "imgs.npy").exists() or np.load(pool / "imgs.npy", mmap_mode="r").shape[0] != args.identities:
|
||
render_pool(pool, args.identities, size, args.workers, args.seed * 104729)
|
||
if args.measure_only:
|
||
# The float network and the generator of an earlier run of this tool, from --out.
|
||
fnet = Net().to(dev)
|
||
fnet.load_state_dict(torch.load(out / "float.pt", map_location=dev))
|
||
fnet.eval()
|
||
blob = (out / "generator.wfg").read_bytes()
|
||
report["generator"] = {"sha256": hashlib.sha256(blob).hexdigest(), "bytes": len(blob)}
|
||
prev = json.loads((out / "report.json").read_text()) if (out / "report.json").exists() else {}
|
||
for k in ("losses", "train_seconds"):
|
||
if k in prev:
|
||
report[k] = prev[k]
|
||
return fnet, blob
|
||
net = Net().to(dev)
|
||
with torch.no_grad():
|
||
# Start with an even mask, and smaller weights.
|
||
net.convs[9].bias[3] = 0.0
|
||
for c in net.convs:
|
||
c.weight.mul_(0.7)
|
||
|
||
from torch.utils.data import DataLoader
|
||
|
||
def batches(base):
|
||
dl = DataLoader(FaceStream(pool, base, gauss, bg), batch_size=args.batch, num_workers=args.workers, persistent_workers=False)
|
||
return iter(dl)
|
||
|
||
def loss_fn(y, t):
|
||
l1 = (y - t).abs().mean()
|
||
gy = lambda a: a[..., 1:, :] - a[..., :-1, :] # noqa: E731
|
||
gx = lambda a: a[..., :, 1:] - a[..., :, :-1] # noqa: E731
|
||
return l1 + 1.0 * ((gy(y) - gy(t)).abs().mean() + (gx(y) - gx(t)).abs().mean())
|
||
|
||
losses = []
|
||
phases = [("float", args.steps, args.lr, False), ("qat", args.qat_steps, args.lr * 0.2, True)]
|
||
it = batches(args.seed * 7919)
|
||
t0 = time.time()
|
||
for name, steps, lr, qat in phases:
|
||
net.qat = qat
|
||
opt = torch.optim.Adam(net.parameters(), lr=lr)
|
||
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, max(steps, 1))
|
||
run = 0.0
|
||
for step in range(steps):
|
||
w_, r_, h_, t_ = (a.to(dev) for a in next(it))
|
||
y, g_, z_ = net(w_, r_, h_)
|
||
# The blend's own loss, the generated picture's (so the generator learns a whole
|
||
# face whatever the mask chooses), and a penalty on mask logits past ±4 (so the
|
||
# mask never saturates where its gradient vanishes).
|
||
loss = loss_fn(y, t_) + 0.25 * (g_ - t_).abs().mean() + 0.1 * torch.relu(z_.abs() - 4).mean()
|
||
opt.zero_grad()
|
||
loss.backward()
|
||
opt.step()
|
||
sched.step()
|
||
run += loss.item()
|
||
if (step + 1) % 500 == 0:
|
||
losses.append([name, step + 1, run / 500])
|
||
print(f"{name} {step + 1}/{steps} loss {run / 500:.3f} ({time.time() - t0:.0f} s)", flush=True)
|
||
run = 0.0
|
||
if name == "float":
|
||
float_state = {k: v.detach().clone() for k, v in net.state_dict().items()}
|
||
torch.save(float_state, out / "float.pt")
|
||
report["losses"] = losses
|
||
report["train_seconds"] = round(time.time() - t0)
|
||
|
||
# Export: each float32 weight times 2^12 and each bias times 2^27, rounded half to even.
|
||
layers = []
|
||
for c in net.convs:
|
||
w = c.weight.detach().cpu().numpy().astype(np.float32).astype(np.float64)
|
||
b = c.bias.detach().cpu().numpy().astype(np.float32).astype(np.float64)
|
||
wi = np.clip(np.round(w * 4096), -32768, 32767).astype(np.int64)
|
||
bi = np.clip(np.round(b * (1 << FB)), -(1 << 31), (1 << 31) - 1).astype(np.int64)
|
||
layers.append((wi, bi))
|
||
hdr = dict(k=K, width=size, height=size, c1=c1, c2=c2, c3=c3, gauss=gauss, bg=bg)
|
||
blob = write_wfg(hdr, layers)
|
||
(out / "generator.wfg").write_bytes(blob)
|
||
report["generator"] = {"sha256": hashlib.sha256(blob).hexdigest(), "bytes": len(blob)}
|
||
|
||
# The float network, for the measurement: the float phase's weights, run in float32.
|
||
fnet = Net().to(dev)
|
||
fnet.load_state_dict(float_state)
|
||
fnet.qat = False
|
||
fnet.eval()
|
||
return fnet, blob
|
||
|
||
|
||
def train_detector(args, dev, pool: Path):
|
||
"""A float keypoint detector, trained on the rendered training faces and their keypoints,
|
||
for the measurement's keypoint error: it reads a picture and gives its sixteen keypoints."""
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
torch.manual_seed(args.seed + 1)
|
||
size = args.size
|
||
|
||
class Det(nn.Module):
|
||
"""Keypoint maps at half the picture's size, each read by a spatial softmax as the
|
||
expected position: sub-pixel by construction."""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
|
||
def block(a, b, s):
|
||
return [nn.Conv2d(a, b, 3, stride=s, padding=1), nn.BatchNorm2d(b), nn.ReLU()]
|
||
self.down = nn.Sequential(*block(3, 32, 1), *block(32, 64, 2), *block(64, 64, 1), *block(64, 128, 2), *block(128, 128, 1), *block(128, 128, 1))
|
||
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False), *block(128, 64, 1), nn.Conv2d(64, K, 3, padding=1))
|
||
half = size // 2
|
||
g = (torch.arange(half, dtype=torch.float32) + 0.5) * 2
|
||
self.register_buffer("gx", g.view(1, 1, 1, half))
|
||
self.register_buffer("gy", g.view(1, 1, half, 1))
|
||
|
||
def forward(self, x):
|
||
m = self.up(self.down(x / 255 - 0.5))
|
||
n, k, h, w = m.shape
|
||
p = torch.softmax(m.reshape(n, k, -1), -1).reshape(n, k, h, w)
|
||
return torch.stack([(p * self.gx).sum((2, 3)), (p * self.gy).sum((2, 3))], -1)
|
||
|
||
det = Det().to(dev)
|
||
opt = torch.optim.Adam(det.parameters(), lr=1e-3)
|
||
det.train()
|
||
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, args.det_steps)
|
||
rng = np.random.default_rng(args.seed + 99)
|
||
all_imgs = np.load(pool / "imgs.npy", mmap_mode="r")
|
||
all_kps = np.load(pool / "kps.npy", mmap_mode="r")
|
||
n = min(args.det_pool, all_imgs.shape[0])
|
||
imgs = torch.from_numpy(np.ascontiguousarray(np.asarray(all_imgs[:n]).reshape(-1, size, size, 3).transpose(0, 3, 1, 2)).astype(np.float32)).to(dev)
|
||
kps = torch.from_numpy(np.asarray(all_kps[:n]).reshape(-1, K, 2).astype(np.float32)).to(dev)
|
||
for step in range(args.det_steps):
|
||
idx = torch.from_numpy(rng.integers(0, imgs.shape[0], args.batch)).to(dev)
|
||
loss = (det(imgs[idx]) - kps[idx]).abs().mean()
|
||
opt.zero_grad()
|
||
loss.backward()
|
||
opt.step()
|
||
sched.step()
|
||
if (step + 1) % 1000 == 0:
|
||
print(f"detector {step + 1}/{args.det_steps} loss {loss.item():.3f} px", flush=True)
|
||
det.eval()
|
||
return det
|
||
|
||
|
||
def psnr(a: np.ndarray, b: np.ndarray) -> float:
|
||
mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
|
||
return float("inf") if mse == 0 else 10 * math.log10(255 * 255 / mse)
|
||
|
||
|
||
def ssim(a: np.ndarray, b: np.ndarray) -> float:
|
||
"""SSIM (11×11 Gaussian window, σ 1.5, K1 0.01, K2 0.03) on each channel, averaged."""
|
||
g = np.exp(-((np.arange(11) - 5) ** 2) / (2 * 1.5 ** 2))
|
||
g /= g.sum()
|
||
|
||
def blur(x):
|
||
x = np.apply_along_axis(lambda r: np.convolve(r, g, mode="valid"), 0, x)
|
||
return np.apply_along_axis(lambda r: np.convolve(r, g, mode="valid"), 1, x)
|
||
|
||
c1, c2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2
|
||
vals = []
|
||
for ch in range(3):
|
||
x, y = a[..., ch].astype(np.float64), b[..., ch].astype(np.float64)
|
||
mx, my = blur(x), blur(y)
|
||
sxx = blur(x * x) - mx * mx
|
||
syy = blur(y * y) - my * my
|
||
sxy = blur(x * y) - mx * my
|
||
s = ((2 * mx * my + c1) * (2 * sxy + c2)) / ((mx * mx + my * my + c1) * (sxx + syy + c2))
|
||
vals.append(s.mean())
|
||
return float(np.mean(vals))
|
||
|
||
|
||
def measure(args, out: Path, report: dict, fnet, blob: bytes):
|
||
import torch
|
||
|
||
dev = next(fnet.parameters()).device
|
||
det = train_detector(args, dev, out / "pool")
|
||
hdr, layers = read_wfg(blob)
|
||
ilayers = [(w, b, s) for (w, b, s) in layers]
|
||
size, gauss, bg = args.size, args.gauss, args.bg
|
||
rows = {"int_float": [], "int_truth": [], "float_truth": [], "warp_truth": []}
|
||
ss = {"int_float": [], "int_truth": [], "float_truth": [], "warp_truth": []}
|
||
kerr = {"truth": [], "int": [], "float": []}
|
||
mouth = {"truth": [], "int": [], "float": []}
|
||
bitrate = {"truth_step4": [], "truth_step1": [], "detected_step4": [], "detected_step1": []}
|
||
gen_sha = hashlib.sha256(blob).digest()
|
||
nseq, nfr = args.measure_identities, args.measure_frames
|
||
t_int = 0.0
|
||
for s in range(nseq):
|
||
ref, rk, frames, kps = sequence(50_000_000 + 7 * s, nfr, size)
|
||
rq = kp_q4(rk)
|
||
ident = hashlib.sha256(b"identity" + struct.pack("<I", s)).digest()
|
||
dq_all = kp_q4(kps)
|
||
for name, src in (("truth", dq_all),):
|
||
for step in (4, 1):
|
||
tot = 0
|
||
for f0 in range(0, nfr, 25):
|
||
tot += len(encode_track(src[f0:f0 + 25].astype(np.float64), step, gen_sha, ident, f0))
|
||
bitrate[f"{name}_step{step}"].append(tot * 8 / (nfr / 25) / 1000)
|
||
with torch.no_grad():
|
||
dk = det(torch.from_numpy(np.ascontiguousarray(frames.transpose(0, 3, 1, 2)).astype(np.float32)).to(dev)).cpu().numpy()
|
||
for step in (4, 1):
|
||
tot = 0
|
||
for f0 in range(0, nfr, 25):
|
||
tot += len(encode_track(kp_q4(dk[f0:f0 + 25]).astype(np.float64), step, gen_sha, ident, f0))
|
||
bitrate[f"detected_step{step}"].append(tot * 8 / (nfr / 25) / 1000)
|
||
ints, floats = [], []
|
||
for i in range(nfr):
|
||
dq = dq_all[i]
|
||
warped, heat = warp_and_maps(ref, rq, dq, gauss, bg)
|
||
t1 = time.time()
|
||
yi = int_forward(ilayers, warped, ref, heat)
|
||
t_int += time.time() - t1
|
||
with torch.no_grad():
|
||
f = lambda a: torch.from_numpy(a.transpose(2, 0, 1)[None].astype(np.float32)).to(dev) # noqa: E731
|
||
yf = fnet(f(warped), f(ref), torch.from_numpy(heat[None].astype(np.float32)).to(dev))[0]
|
||
yf = np.clip(np.floor(yf[0].cpu().numpy().transpose(1, 2, 0) + 0.5), 0, 255).astype(np.uint8)
|
||
truth = frames[i]
|
||
for key, (a, b) in {"int_float": (yi, yf), "int_truth": (yi, truth), "float_truth": (yf, truth), "warp_truth": (warped, truth)}.items():
|
||
rows[key].append(psnr(a, b))
|
||
ss[key].append(ssim(a, b))
|
||
ints.append(yi)
|
||
floats.append(yf)
|
||
with torch.no_grad():
|
||
for key, imgs in (("int", np.stack(ints)), ("float", np.stack(floats)), ("truth", frames)):
|
||
d = det(torch.from_numpy(np.ascontiguousarray(imgs.transpose(0, 3, 1, 2)).astype(np.float32)).to(dev)).cpu().numpy()
|
||
err = np.sqrt(((d - kps) ** 2).sum(-1))
|
||
kerr[key].append(float(np.sqrt((err ** 2).mean())))
|
||
mouth[key].append(float(np.sqrt((err[:, 7:11] ** 2).mean())))
|
||
print(f"measured sequence {s + 1}/{nseq}", flush=True)
|
||
finite = lambda v: [x for x in v if math.isfinite(x)] # noqa: E731
|
||
report["measurement"] = {
|
||
"identities": nseq, "frames_each": nfr, "fps": 25, "size": [size, size],
|
||
"psnr_db": {k: {"mean": round(float(np.mean(finite(v))), 2), "min": round(float(np.min(v)), 2), "identical_frames": int(sum(1 for x in v if not math.isfinite(x)))} for k, v in rows.items()},
|
||
"ssim": {k: {"mean": round(float(np.mean(v)), 4), "min": round(float(np.min(v)), 4)} for k, v in ss.items()},
|
||
"keypoint_rms_px": {k: round(float(np.mean(v)), 3) for k, v in kerr.items()},
|
||
"mouth_keypoint_rms_px": {k: round(float(np.mean(v)), 3) for k, v in mouth.items()},
|
||
"track_kbit_s": {k: round(float(np.mean(v)), 3) for k, v in bitrate.items()},
|
||
"numpy_int_seconds_per_frame": round(t_int / (nseq * nfr), 4),
|
||
}
|
||
|
||
|
||
def corpus_inputs(args, out: Path, report: dict, blob: bytes):
|
||
"""Identity priors and keypoint tracks of held-out faces, for the corpus (§9)."""
|
||
gen = hashlib.sha256(blob).hexdigest()
|
||
(out / "identities").mkdir(exist_ok=True)
|
||
(out / "tracks").mkdir(exist_ok=True)
|
||
files = {}
|
||
for name, seed, nfr in (("face1", 70_000_001, args.corpus_frames), ("face2", 70_000_017, args.corpus_frames)):
|
||
ref, rk, frames, kps = sequence(seed, nfr, args.size)
|
||
rq = kp_q4(rk)
|
||
wfi = bytearray(b"WFI1")
|
||
wfi += struct.pack("<BBHH", 1, K, args.size, args.size)
|
||
wfi += rq.astype("<i2").tobytes()
|
||
wfi += ref.tobytes()
|
||
(out / "identities" / f"{name}.wfi").write_bytes(bytes(wfi))
|
||
track = {"identity": name, "keypoints": kp_q4(kps).tolist()}
|
||
(out / "tracks" / f"{name}.json").write_text(json.dumps(track, separators=(",", ":")) + "\n")
|
||
files[name] = {"wfi_sha256": hashlib.sha256(bytes(wfi)).hexdigest(), "frames": nfr}
|
||
report["corpus_inputs"] = {"generator": gen, "faces": files}
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
|
||
ap.add_argument("--out", type=Path, required=True)
|
||
ap.add_argument("--seed", type=int, default=20261022)
|
||
ap.add_argument("--size", type=int, default=64)
|
||
ap.add_argument("--c1", type=int, default=24)
|
||
ap.add_argument("--c2", type=int, default=48)
|
||
ap.add_argument("--c3", type=int, default=64)
|
||
ap.add_argument("--gauss", type=int, default=1024, help="§7.5's g: 32768/σ², σ in pixels")
|
||
ap.add_argument("--bg", type=int, default=6554, help="§7.5's background weight, Q16")
|
||
ap.add_argument("--identities", type=int, default=6000)
|
||
ap.add_argument("--steps", type=int, default=60000)
|
||
ap.add_argument("--qat-steps", type=int, default=12000)
|
||
ap.add_argument("--lr", type=float, default=1e-3)
|
||
ap.add_argument("--batch", type=int, default=32)
|
||
ap.add_argument("--workers", type=int, default=12)
|
||
ap.add_argument("--det-steps", type=int, default=12000)
|
||
ap.add_argument("--det-pool", type=int, default=6000, help="training identities the detector reads, all their poses")
|
||
ap.add_argument("--measure-identities", type=int, default=12)
|
||
ap.add_argument("--measure-frames", type=int, default=100)
|
||
ap.add_argument("--corpus-frames", type=int, default=3)
|
||
ap.add_argument("--preview", action="store_true", help="only write preview pictures of the renderer and the warp")
|
||
ap.add_argument("--measure-only", action="store_true", help="measure the generator and float network an earlier run wrote to --out")
|
||
args = ap.parse_args()
|
||
args.out.mkdir(parents=True, exist_ok=True)
|
||
if args.preview:
|
||
preview(args)
|
||
return
|
||
report = {"seed": args.seed, "args": {k: v for k, v in vars(args).items() if k != "out"}, "keypoints": KEYPOINTS}
|
||
fnet, blob = train(args, args.out, report)
|
||
corpus_inputs(args, args.out, report, blob)
|
||
measure(args, args.out, report, fnet, blob)
|
||
(args.out / "report.json").write_text(json.dumps(report, indent=1, sort_keys=True) + "\n")
|
||
print(json.dumps(report["measurement"], indent=1))
|
||
|
||
|
||
def preview(args):
|
||
"""A contact sheet of generated faces and warps, as a binary PPM (no image library)."""
|
||
tiles = []
|
||
for i in range(8):
|
||
ref, warped, heat, tgt, rq, tq = make_sample(1000 + i, args.size, args.gauss, args.bg)
|
||
hm = np.clip(heat.max(0) * 255 // 256, 0, 255).astype(np.uint8)
|
||
tiles.append(np.concatenate([ref, warped, tgt, np.stack([hm] * 3, -1)], 1))
|
||
sheet = np.concatenate(tiles, 0)
|
||
big = np.repeat(np.repeat(sheet, 3, 0), 3, 1)
|
||
h, w, _ = big.shape
|
||
(args.out / "preview.ppm").write_bytes(b"P6 %d %d 255\n" % (w, h) + big.tobytes())
|
||
_, _, frames, kps = sequence(50_000_000, 16, args.size)
|
||
seq = np.concatenate(list(frames), 1)
|
||
seq = np.repeat(np.repeat(seq, 3, 0), 3, 1)
|
||
h, w, _ = seq.shape
|
||
(args.out / "sequence.ppm").write_bytes(b"P6 %d %d 255\n" % (w, h) + seq.tobytes())
|
||
print("wrote", args.out / "preview.ppm", args.out / "sequence.ppm")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|