Dragon-s-Lair-X68k/tools/encoder/vq_hybrid.py

#!/usr/bin/env python3
"""Cinepak-style hybrid VQ with a RAW escape: per-4x4-block choice of
SKIP / V1 (one 4x4 codeword) / V4 (four 2x2 codewords) / RAW (16 literal indices).

Flat 4x4 VQ at k=256 visibly destroys Bluth's ink linework (see docs/FINDINGS.md).
The standard fix is to let detailed blocks spend 4x the bits. Rate control picks
the split per block by rate-distortion, so the bitrate ceiling stays deterministic
-- which is the whole reason we chose VQ over a lossless delta.

The RAW mode is what makes ONE codec serve both shipping targets (session 2
user decision: SASI and SCSI quality modes). As lam -> 0 the encoder buys RAW
blocks until the frame is pixel-exact against the palettised source, so the
SCSI profile is not a second codec -- it is the same bitstream with the rate
knob opened up. The 68000 decoder needs no extra path: RAW is a straight copy,
which is cheaper than V4.

Bitstream per frame (what the 68000 actually parses):
  2 bits/block header, packed:  00=SKIP  01=V1  10=V4  11=RAW
  then the payload in block order: V1 -> 1 index, V4 -> 4, RAW -> 16

STRUCTURE (session 6). The encoder is FRAME-DRIVABLE: `frame_ctx` / `decide` /
`paint` expose one frame at a time so a caller can choose `lam` per frame and
feed back the frame it actually emitted. That is not a convenience -- it is the
fix for FINDINGS 26.1. This codec is temporally recursive (SKIP copies the
previous RECONSTRUCTION), so any rate control that picks frames out of
independently-encoded whole-sequence runs desynchronises the encoder from the
decoder. `encode()` is now a thin loop over the per-frame API and stays the
fixed-lam path.

The split is also what makes rate control affordable: `l1`/`l4g` and their
errors depend on neither `lam` nor `prev`, so they are computed once per frame
and a lam search only re-runs the argmin.
"""
import numpy as np, sys
from PIL import Image
import vq as VQ

LUMA = VQ.LUMA

# byte cost of each mode's payload, per block. The 2-bit header is paid by
# every block regardless, so it drops out of the mode comparison.
_HDR_BYTES_PER_BLOCK = 2 / 8.0
RAW_BYTES = 16.0                  # literal palette bytes, never indices


def blocks_of(idx, pal, bw, bh):
    return VQ.blockify(idx, pal, bw, bh)


def build(frames_dir, k1=256, k4=256, iters=16, lam=0.0):
    rgb = VQ.load_frames(frames_dir)
    H, W = rgb[0].shape[:2]
    ref, pal = VQ.scene_palette(rgb)
    idx = VQ.palettise(rgb, ref)

    # --- two codebooks, trained on the whole scene ---
    X1 = np.concatenate([blocks_of(i, pal, 4, 4) for i in idx])
    C1, _ = VQ.kmeans(X1, k1, iters)
    cb1 = VQ.snap_codebook(C1, pal, 4, 4)                  # (k1,16) palette idx
    C1s = (pal[cb1].astype(np.float32) * LUMA).reshape(k1, -1)

    X4 = np.concatenate([blocks_of(i, pal, 2, 2) for i in idx])
    C4, _ = VQ.kmeans(X4, k4, iters)
    cb4 = VQ.snap_codebook(C4, pal, 2, 2)                  # (k4,4) palette idx
    C4s = (pal[cb4].astype(np.float32) * LUMA).reshape(k4, -1)
    return dict(rgb=rgb, pal=pal, idx=idx, H=H, W=W,
                cb1=cb1, C1s=C1s, cb4=cb4, C4s=C4s, k1=k1, k4=k4,
                nbx=W // 4, nby=H // 4, nb=(W // 4) * (H // 4),
                palw=pal.astype(np.float32) * LUMA)


def _v1_recon(lab1, cb1, H, W):
    return VQ.unblockify(lab1, cb1, H, W, 4, 4)


# --- block <-> image reshapes (no copy where numpy can avoid one) -------------

def to_blocks(a, nbx, nby):
    """(H,W) -> (nb,4,4) in block raster order."""
    return a.reshape(nby, 4, nbx, 4).transpose(0, 2, 1, 3).reshape(-1, 4, 4)


def from_blocks(b, nbx, nby):
    """(nb,4,4) -> (H,W)."""
    return b.reshape(nby, nbx, 4, 4).transpose(0, 2, 1, 3).reshape(nby * 4, nbx * 4)


def default_idx_bytes(m):
    """Size of ONE codebook index in the bitstream. k>256 needs 2 bytes, which
    doubles what V1 and V4 actually cost -- an RD model that ignores that
    systematically over-picks V4 and under-reports the bitrate (FINDINGS 14)."""
    return 1 if max(m["k1"], m["k4"]) <= 256 else 2


# --- per-frame API -----------------------------------------------------------

def frame_symbols(m, f):
    """lam- and prev-INDEPENDENT part of a frame: codeword assignments and
    their errors.

    Cached, because a lam search re-uses them unchanged and `VQ.assign` is the
    expensive call in the encoder -- 22.8 of 24.6 ms per frame, measured. That
    cache is what makes per-frame rate control affordable: a 12-step search
    over 120 frames costs 0.3 s, against 49 s for the equivalent done by
    re-running whole-sequence encodes.

    The cache holds ONE frame. Every caller works a frame at a time, and at
    ~133 KB of intermediates per frame a whole-sequence cache would cost
    900 MB on a 9.4-minute stream for no benefit."""
    cache = m.get("_sym")
    if cache is not None and cache[0] == f:
        return cache[1]
    pal, W, nb = m["pal"], m["W"], m["nb"]
    im = m["idx"][f]

    B1 = blocks_of(im, pal, 4, 4)                       # (nb,48)
    l1 = VQ.assign(B1, m["C1s"])
    e1 = ((B1 - m["C1s"][l1]) ** 2).sum(1)

    B4 = blocks_of(im, pal, 2, 2)                       # (nb*4,12) in 2x2 raster
    l4 = VQ.assign(B4, m["C4s"])
    e4raw = ((B4 - m["C4s"][l4]) ** 2).sum(1)
    # regroup 2x2 blocks into their parent 4x4 block
    q = _group_2x2_into_4x4(np.arange(nb * 4), W)
    e4 = e4raw[q].reshape(nb, 4).sum(1)
    l4g = l4[q].reshape(nb, 4)

    s = dict(l1=l1, e1=e1, l4g=l4g, e4=e4,
             src_blocks=to_blocks(im, m["nbx"], m["nby"]))
    m["_sym"] = (f, s)
    return s


def frame_ctx(m, f, prev, idx_bytes=None):
    """Everything needed to decide one frame at any lam, given the frame that
    will actually precede it in the emitted stream."""
    s = frame_symbols(m, f)
    nb = m["nb"]
    if prev is None:
        eS = np.full(nb, np.inf)
        prev_blocks = None
    else:
        # SKIP distortion = this frame against the previous RECONSTRUCTION,
        # in the same luma-weighted space the codebooks were trained in.
        d = ((m["palw"][m["idx"][f]] - m["palw"][prev]) ** 2).sum(2)
        eS = d.reshape(m["nby"], 4, m["nbx"], 4).sum((1, 3)).ravel()
        prev_blocks = to_blocks(prev, m["nbx"], m["nby"])
    return dict(f=f, sym=s, eS=eS, prev_blocks=prev_blocks, nb=nb,
                idx_bytes=default_idx_bytes(m) if idx_bytes is None else idx_bytes)


def decide(ctx, lam):
    """Lagrangian mode decision at one lam. Returns (mode, payload bytes).

    Cheap by design: no painting, no image-sized work. A lam search calls this
    a dozen times per frame and paints once."""
    ib = ctx["idx_bytes"]
    s = ctx["sym"]
    cost = np.stack([ctx["eS"],
                     s["e1"] + lam * (1.0 * ib),
                     s["e4"] + lam * (4.0 * ib),
                     np.full(ctx["nb"], lam * RAW_BYTES)])
    mode = np.argmin(cost, axis=0).astype(np.uint8)
    return mode, frame_bytes(mode, ctx["nb"], ib)


def frame_bytes(mode, nb, idx_bytes):
    nV1 = int((mode == 1).sum()); nV4 = int((mode == 2).sum())
    nR = int((mode == 3).sum())
    return (nb * _HDR_BYTES_PER_BLOCK
            + (nV1 + nV4 * 4) * idx_bytes + nR * RAW_BYTES)


def paint(m, ctx, mode):
    """Reconstruct the frame the decoder will produce for this mode map."""
    nbx, nby = m["nbx"], m["nby"]
    s = ctx["sym"]
    ob = np.empty((ctx["nb"], 4, 4), dtype=np.uint8)
    sel = mode == 0
    if sel.any():
        ob[sel] = ctx["prev_blocks"][sel]
    sel = mode == 1
    if sel.any():
        ob[sel] = m["cb1"][s["l1"][sel]].reshape(-1, 4, 4)
    sel = mode == 2
    if sel.any():
        # (n, sub_y, sub_x, py, px) -> (n, sub_y, py, sub_x, px) -> (n,4,4)
        c = m["cb4"][s["l4g"][sel]].reshape(-1, 2, 2, 2, 2)
        ob[sel] = c.transpose(0, 1, 3, 2, 4).reshape(-1, 4, 4)
    sel = mode == 3
    if sel.any():
        ob[sel] = s["src_blocks"][sel]
    return from_blocks(ob, nbx, nby)


def encode_frame(m, f, prev, lam, idx_bytes=None):
    """One frame at one lam against one previous reconstruction."""
    ctx = frame_ctx(m, f, prev, idx_bytes)
    mode, sz = decide(ctx, lam)
    return dict(recon=paint(m, ctx, mode), mode=mode, size=sz,
                l1=ctx["sym"]["l1"], l4g=ctx["sym"]["l4g"], ctx=ctx)


def encode(m, lam=0.02, skip_thresh=0.0, idx_bytes=None):
    """Fixed-lam whole-sequence encode: a loop over the per-frame API.

    lam = lagrangian rate weight (bytes -> squared-error units).
    Higher lam => more V1/SKIP => smaller & softer.

    For a rate-controlled encode use ratectl.encode_rate_controlled(), which
    drives the same per-frame API and varies lam. Do NOT reassemble a sequence
    out of several fixed-lam runs of this function -- FINDINGS 26.1."""
    if idx_bytes is None:
        idx_bytes = default_idx_bytes(m)
    recon, modes, sizes, l1s, l4gs = [], [], [], [], []
    prev = None
    for f in range(len(m["idx"])):
        r = encode_frame(m, f, prev, lam, idx_bytes)
        recon.append(r["recon"]); modes.append(r["mode"]); sizes.append(r["size"])
        l1s.append(r["l1"]); l4gs.append(r["l4g"])
        prev = r["recon"]
    return dict(recon=recon, modes=modes, sizes=np.array(sizes),
                l1=l1s, l4g=l4gs, nb=m["nb"])


def _group_2x2_into_4x4(a, W):
    """map 2x2-block raster order -> (nb4, 4) grouping by parent 4x4 block"""
    n2x = W // 2
    n2y = len(a) // n2x
    g = a.reshape(n2y, n2x)
    g = g.reshape(n2y // 2, 2, n2x // 2, 2).transpose(0, 2, 1, 3)
    return g.reshape(-1)


def evaluate(m, enc, fps=12):
    pal = m["pal"]
    rec = [pal[i] for i in enc["recon"]]
    src = [pal[i] for i in m["idx"]]
    p_vq = np.mean([VQ.psnr(o, v) for o, v in zip(m["rgb"], rec)])
    p_pal = np.mean([VQ.psnr(o, v) for o, v in zip(m["rgb"], src)])
    mo = np.concatenate(enc["modes"])
    sz = enc["sizes"].mean()
    return dict(psnr=p_vq, pal=p_pal, loss=p_pal - p_vq, bytes=sz,
                kbps=sz * fps / 1024,
                skip=100 * (mo == 0).mean(), v1=100 * (mo == 1).mean(),
                v4=100 * (mo == 2).mean(), raw=100 * (mo == 3).mean())