notes.ell.dev
All notesindex.org

encoder.py

python
# /// script
# dependencies = ["numpy", "pillow"]
# ///
"""Encode a video clip into a self-contained GLSL shader using BC1-style 4x4 blocks.

Each 4x4 block = two RGB565 endpoints (one uint) + sixteen 2-bit lerp indices (one uint),
so any pixel of any frame decodes with two array reads. 4 bits per pixel.
"""
import argparse, subprocess, sys
import numpy as np
from PIL import Image

ap = argparse.ArgumentParser()
ap.add_argument("src"); ap.add_argument("out")
ap.add_argument("--width", type=int, default=128)
ap.add_argument("--height", type=int, default=72)
ap.add_argument("--fps", type=float, default=10)
ap.add_argument("--seconds", type=float, default=3)
ap.add_argument("--start", type=float, default=0)
ap.add_argument("--preview", help="write a decoded contact sheet PNG here")
ap.add_argument("--name", default="clip")
a = ap.parse_args()

W, H = a.width, a.height
assert W % 4 == 0, "width must be a multiple of 4"
HP = (H + 3) // 4 * 4
raw = subprocess.run(
    ["ffmpeg", "-v", "error", "-ss", str(a.start), "-t", str(a.seconds), "-i", a.src,
     "-vf", f"fps={a.fps},scale={W}:{H}:flags=area", "-f", "rawvideo", "-pix_fmt", "rgb24", "-"],
    check=True, capture_output=True).stdout
frames = np.frombuffer(raw, np.uint8).reshape(-1, H, W, 3).astype(np.float32) / 255.0
if HP != H:  # replicate the last row so edge blocks fit cleanly
    frames = np.concatenate([frames, np.repeat(frames[:, -1:], HP - H, axis=1)], axis=1)
F = frames.shape[0]
BW, BH = W // 4, HP // 4

# (F, BH, BW, 16, 3) blocks, pixel order y*4+x
blk = frames.reshape(F, BH, 4, BW, 4, 3).transpose(0, 1, 3, 2, 4, 5).reshape(-1, 16, 3)

def q565(c):
    c = np.clip(c, 0, 1)
    r = np.round(c[..., 0] * 31).astype(np.uint32)
    g = np.round(c[..., 1] * 63).astype(np.uint32)
    b = np.round(c[..., 2] * 31).astype(np.uint32)
    return (r << 11) | (g << 5) | b

def d565(v):
    return np.stack([((v >> 11) & 31) / 31, ((v >> 5) & 63) / 63, (v & 31) / 31], -1).astype(np.float32)

W4 = np.array([0, 1 / 3, 2 / 3, 1], np.float32)

def assign(p, e0, e1):
    pal = e0[:, None] + (e1 - e0)[:, None] * W4[None, :, None]           # (N,4,3)
    d = ((p[:, :, None] - pal[:, None]) ** 2).sum(-1)                     # (N,16,4)
    idx = d.argmin(-1)
    return idx, np.take_along_axis(d, idx[..., None], -1)[..., 0].sum(-1)

# Initial endpoints: extremes along the principal axis (power iteration on covariance).
mean = blk.mean(1)
cen = blk - mean[:, None]
cov = np.einsum("npi,npj->nij", cen, cen)
ax = np.ones((len(blk), 3), np.float32)
for _ in range(8):
    ax = np.einsum("nij,nj->ni", cov, ax)
    ax /= np.linalg.norm(ax, axis=1, keepdims=True) + 1e-12
t = np.einsum("npi,ni->np", cen, ax)
c0 = q565(mean + t.min(1)[:, None] * ax); c1 = q565(mean + t.max(1)[:, None] * ax)
idx, err = assign(blk, d565(c0), d565(c1))

# Refine: least-squares endpoints for the chosen weights, keep whichever is better.
for _ in range(2):
    w = W4[idx]                                                          # (N,16)
    a11 = ((1 - w) ** 2).sum(1); a22 = (w ** 2).sum(1); a12 = ((1 - w) * w).sum(1)
    det = a11 * a22 - a12 ** 2
    ok = np.abs(det) > 1e-8
    b1 = ((1 - w)[..., None] * blk).sum(1); b2 = (w[..., None] * blk).sum(1)
    s = np.where(ok, 1 / np.where(ok, det, 1), 0)[:, None]
    n0 = q565((a22[:, None] * b1 - a12[:, None] * b2) * s)
    n1 = q565((a11[:, None] * b2 - a12[:, None] * b1) * s)
    nidx, nerr = assign(blk, d565(n0), d565(n1))
    better = ok & (nerr < err)
    c0 = np.where(better, n0, c0); c1 = np.where(better, n1, c1)
    idx = np.where(better[:, None], nidx, idx); err = np.where(better, nerr, err)

bits = (idx.astype(np.uint32) << (2 * np.arange(16, dtype=np.uint32))).sum(1).astype(np.uint32)
data = np.stack([c0 | (c1 << 16), bits], 1).reshape(-1)

# Decode back for an error report and an optional contact sheet.
dec = d565(c0)[:, None] + (d565(c1) - d565(c0))[:, None] * W4[idx][..., None]
dec = dec.reshape(F, BH, BW, 4, 4, 3).transpose(0, 1, 3, 2, 4, 5).reshape(F, HP, W, 3)[:, :H]
ref = frames[:, :H]
psnr = 10 * np.log10(1 / ((dec - ref) ** 2).mean())
if a.preview:
    pick = np.linspace(0, F - 1, 8).astype(int)
    rows = [np.concatenate([np.concatenate(list(ref[pick[i:i + 4]]), 1),
                            ], 0) for i in (0, 4)]
    top = np.concatenate([np.concatenate(list(ref[pick[:4]]), 1), np.concatenate(list(dec[pick[:4]]), 1)], 0)
    bot = np.concatenate([np.concatenate(list(ref[pick[4:]]), 1), np.concatenate(list(dec[pick[4:]]), 1)], 0)
    img = Image.fromarray((np.concatenate([top, bot], 0) * 255).round().astype(np.uint8))
    img.resize((img.width * 4, img.height * 4), Image.NEAREST).save(a.preview)

# Store the data as a binary tree of ifs returning uvec4 constants (two blocks per leaf). On NVIDIA's GL
# driver, dynamically indexed const arrays are capped at ~32 KB per shader, and a switch with more than a few
# hundred cases silently returns zeros. The tree reads back correctly and is limited by the ~65k
# instruction cap (about 128 KB of data).
quads = np.concatenate([data, np.zeros((-len(data)) % 4, np.uint32)]).reshape(-1, 4).tolist()
def tree(lo, hi):
    if hi - lo == 1:
        return "return uvec4(" + ",".join(f"0x{v:x}u" for v in quads[lo]) + ");\n"
    m = (lo + hi) // 2
    return f"if (i < {m}) {{\n{tree(lo, m)}}} else {{\n{tree(m, hi)}}}\n"
cases = tree(0, len(quads))

src = f"""// CLIP / {a.name}: {a.seconds:g} s of video baked into the shader, looping inside the "Colonq's Face" window
// (bottom right, same placement as dvd-window.glsl). Everything outside the window is transparent.
// Encoding: {W}x{H} px, {a.fps:g} fps, {F} frames, BC1-style 4x4 blocks (two RGB565 endpoints + 2-bit
// indices per pixel = 4 bits per pixel, {len(quads)} uvec4s). Generated by encode.py; edit the knobs, not the data.
// No interactive inputs: chat_time, tracking_*, cursors and heartrate are unused.
// Paste this entire file into the host; the host supplies the documented uniforms.

// GLES hosts default ints to mediump, which Mesa on AMD (Steam Deck) can run as 16-bit; the decoder needs 32.
precision highp float;
precision highp int;

// ---- Layout, as fractions of the screen (same as dvd-window.glsl).
const vec2 S_WIN_MIN = vec2(0.5985, 0.300);   // Face window, top left.
const vec2 S_WIN_MAX = vec2(1.000, 0.951);    // Face window, bottom right (above its mode line).

// ---- Playback.
const float V_SPEED = 1.0;                    // Playback rate; 0.5 = half speed.
const float V_FILL = 1.0;                     // 1 = fit the window; larger crops in.
const vec3 V_BG = vec3(0.0);                  // Letterbox colour.

// ---- Generated data.
const int V_W = {W};
const int V_H = {H};
const int V_FRAMES = {F};
const float V_FPS = {float(a.fps)!r};
const int V_BW = {BW};
const int V_BLOCKS = {BW * BH};

highp uvec4 V_DATA(highp int i) {{
{cases}}}

vec3 unpack565(uint v) {{
    return vec3(float((v >> 11u) & 31u) / 31.0, float((v >> 5u) & 63u) / 63.0, float(v & 31u) / 31.0);
}}

vec3 texel(highp ivec2 p, highp int f) {{
    p = clamp(p, ivec2(0), ivec2(V_W - 1, V_H - 1));
    highp int b = f * V_BLOCKS + (p.y >> 2) * V_BW + (p.x >> 2);
    highp uvec4 q = V_DATA(b >> 1);
    highp uint ends = q.x, bits = q.y;
    if ((b & 1) == 1) {{
        ends = q.z;
        bits = q.w;
    }}
    highp uint i = (bits >> uint(2 * (((p.y & 3) << 2) | (p.x & 3)))) & 3u;
    return mix(unpack565(ends & 0xffffu), unpack565(ends >> 16u), float(i) / 3.0);
}}

vec4 shade(vec2 cs) {{
    vec2 px = cs * resolution;
    vec2 wmin = S_WIN_MIN * resolution, wmax = S_WIN_MAX * resolution;
    if (any(lessThan(px, wmin)) || any(greaterThan(px, wmax))) return vec4(0.0);

    // Fit the clip into the window, preserving its aspect ratio.
    vec2 wsize = wmax - wmin;
    float scale = min(wsize.x / float(V_W), wsize.y / float(V_H)) * V_FILL;
    vec2 uv = (px - (wmin + wmax) * 0.5) / scale + vec2(V_W, V_H) * 0.5;
    if (any(lessThan(uv, vec2(0.0))) || any(greaterThanEqual(uv, vec2(V_W, V_H)))) return vec4(V_BG, 1.0);

    int f = int(mod(max(time, 0.0) * V_SPEED * V_FPS, float(V_FRAMES)));
    f = clamp(f, 0, V_FRAMES - 1);
    // Exactly one texel() call site: the driver inlines the data switch once per call, so bilinear
    // filtering (four calls) would quadruple compile time and instruction count.
    vec3 c = texel(ivec2(uv), f);
    return vec4(c, 1.0);
}}
"""
open(a.out, "w").write(src)
print(f"{a.out}: {W}x{H} @ {a.fps:g} fps, {F} frames, {len(data)} uints, "
      f"{len(quads)} uvec4s, {len(src) / 1024:.0f} KiB source, {data.nbytes / 1024:.0f} KiB data, PSNR {psnr:.1f} dB")