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
"""
import numpy as np, sys
from PIL import Image
import vq as VQ

LUMA = VQ.LUMA


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)


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


def encode(m, lam=0.02, skip_thresh=0.0, idx_bytes=None):
    """lam = lagrangian rate weight (bytes -> squared-error units).
    Higher lam => more V1/SKIP => smaller & softer.

    idx_bytes: size of ONE codebook index in the bitstream. k>256 needs 2 bytes,
    which doubles what V1 and V4 actually cost -- if the RD model ignores that
    it systematically over-picks V4 and under-reports the bitrate. Defaults to
    the value implied by the codebook sizes."""
    if idx_bytes is None:
        idx_bytes = 1 if max(m["k1"], m["k4"]) <= 256 else 2
    pal, idx, H, W = m["pal"], m["idx"], m["H"], m["W"]
    nbx, nby = W // 4, H // 4
    nb = nbx * nby
    recon, modes, sizes = [], [], []
    prev = None
    for f, im in enumerate(idx):
        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 (raster over 8x12... ) 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)

        # SKIP: cost of reusing the previous *reconstructed* block
        if prev is None:
            eS = np.full(nb, np.inf)
        else:
            pb = blocks_of(prev, pal, 4, 4)
            eS = ((B1 - pb) ** 2).sum(1)

        # RAW: zero distortion against the palettised source, 16 bytes
        eR = np.zeros(nb)

        # rate-distortion choice: true byte cost per mode. The 2-bit header is
        # paid by every block regardless, so it drops out of the comparison.
        bV1 = 1.0 * idx_bytes
        bV4 = 4.0 * idx_bytes
        bRAW = 16.0                       # RAW is literal palette bytes, never indices
        cost = np.stack([eS + lam * 0.0, e1 + lam * bV1,
                         e4 + lam * bV4, eR + lam * bRAW])
        mode = np.argmin(cost, axis=0).astype(np.uint8)

        out = np.empty((H, W), dtype=np.uint8)
        _paint(out, mode, l1, l4g, m["cb1"], m["cb4"], prev, nbx, nby, im)
        recon.append(out); modes.append(mode)
        nV1 = int((mode == 1).sum()); nV4 = int((mode == 2).sum())
        nR = int((mode == 3).sum())
        sizes.append(nb * 2 / 8 + (nV1 + nV4 * 4) * idx_bytes + nR * 16)
        prev = out
    return dict(recon=recon, modes=modes, sizes=np.array(sizes), nb=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 _paint(out, mode, l1, l4g, cb1, cb4, prev, nbx, nby, src):
    for b in range(len(mode)):
        by, bx = divmod(b, nbx)
        y, x = by * 4, bx * 4
        mo = mode[b]
        if mo == 0:
            out[y:y+4, x:x+4] = prev[y:y+4, x:x+4]
        elif mo == 1:
            out[y:y+4, x:x+4] = cb1[l1[b]].reshape(4, 4)
        elif mo == 3:
            out[y:y+4, x:x+4] = src[y:y+4, x:x+4]
        else:
            c = cb4[l4g[b]].reshape(2, 2, 2, 2)             # (sub_y,sub_x,2,2)
            out[y:y+2, x:x+2] = c[0, 0]; out[y:y+2, x+2:x+4] = c[0, 1]
            out[y+2:y+4, x:x+2] = c[1, 0]; out[y+2:y+4, x+2:x+4] = c[1, 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())