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

#!/usr/bin/env python3
"""4x4 vector-quantisation prototype for the X68000 Dragon's Lair port.

Pipeline mirrors what the 68000 decoder would actually do, so the measured
quality is honest:

  frames -> per-scene 256-colour palette (median cut, NO dither)
         -> 4x4 blocks of PALETTISED rgb
         -> k-means codebook (luma-weighted euclidean)
         -> each codeword's 16 pixels snapped back to a palette index

The decoder only ever copies 16 palette indices out of a table, so the codebook
entries MUST be legal palette indices -- both quantisation losses compose.

No sklearn on this box; k-means is hand-rolled (chunked, numpy).
"""
import numpy as np, glob, os, sys
from concurrent.futures import ThreadPoolExecutor
from PIL import Image

BW = BH = 4                      # block size
# ITU-R BT.601 luma weights, squared -- we compare in a luma-weighted RGB space
LUMA = np.array([0.299, 0.587, 0.114], dtype=np.float32)


def load_frames(d):
    fs = sorted(glob.glob(f"{d}/f*.png"))
    return [np.asarray(Image.open(f).convert("RGB")) for f in fs]


def scene_palette(rgb, colors=256, stride=3, reserve_black=True):
    """One shared palette for the whole scene, no dithering (cel art is flat).

    `reserve_black` puts TRUE BLACK at index 0 and quantises the picture into
    the other 255 entries.  It is not a cosmetic default (FINDINGS 23.4): the
    picture is 256x192 inside a 256x256 mode, GVRAM cleared to zero displays
    palette entry 0, and a free mediancut palette puts a real image colour
    there -- on 00020 f0001 it was (206,192,176), used by 210 image pixels, so
    the 64 blank rows of letterbox came out beige.  Entry 0 also needs `I = 0`
    in the X68000's GRB555 word or the bars sit at RGB (4,4,4) (23.3); that
    half is `dlxload.pack_palette`'s and it needs no special case, because a
    (0,0,0) entry picks I=0 by its own minimum-squared-error rule.

    Black is RESERVED, not withheld: the mapper may still spend index 0 on
    genuinely black pixels, which is the entry it would have wanted anyway.
    What the reservation buys is that index 0 is black REGARDLESS of what the
    scene contains, which is what the letterbox needs and what a free palette
    cannot promise.
    """
    samp = np.concatenate([r.reshape(-1, 3) for r in rgb[::stride]])
    n = colors - 1 if reserve_black else colors
    q = Image.fromarray(samp.reshape(-1, 1, 3)).quantize(
        colors=n, method=Image.MEDIANCUT, dither=Image.NONE)
    pal = np.array(q.getpalette()[:n * 3], dtype=np.uint8).reshape(-1, 3)
    if not reserve_black:
        return q, pal
    pal = np.vstack([np.zeros((1, 3), np.uint8), pal])
    # The quantiser above cannot be reused as the mapping reference: its
    # palette is the 255 it chose, at the wrong indices.  A P-mode image
    # carrying the FINAL palette is what every frame is then mapped against,
    # so the indices in the container and the entries in the container's
    # palette section are the same table by construction.
    ref = Image.new("P", (1, 1))
    ref.putpalette(pal.tobytes().ljust(768, b"\0"))
    return ref, pal


def frame_palette(rgb1, colors=254):
    """`scene_palette`'s sibling, for the PACKED layout: ONE FRAME, 254 colours.

    The codec cannot have this.  Every codeword it emits is an index INTO
    `scene_palette`, so its palette is shared scene-wide and 31.33 dB is a
    ceiling no bitrate crosses (FINDINGS 61.9).  A literal frame has no
    codebooks, so nothing forces a shared palette on it.

    The layout spends TWO entries where `--reserve-black` spends one (47.2):
    index 0 is the TRANSPARENCY KEY of the top graphics page and must never
    appear in the picture, and black therefore lives at 255 for the letterbox.
    So the picture gets 254.

    The +1 shift is the whole mechanism, and it is why this does NOT go through
    a P-mode reference image the way `scene_palette` does.  `palettise` maps
    against the FINAL 256-entry table, and that table has (0,0,0) at both 0 and
    255 -- a nearest-colour mapper is free to pick either, and there is no way to
    forbid the one that must stay unused.  Quantising to 254 and shifting keeps
    index 0 free BY CONSTRUCTION rather than by hoping the mapper agrees, and it
    is still exact: `pal[idx]` reproduces the quantiser's own rendering.

    Returns (pal (256,3) uint8, idx (H,W) uint8 in 1..254).
    """
    q = Image.fromarray(rgb1).quantize(colors=colors, method=Image.MEDIANCUT,
                                       dither=Image.NONE)
    raw = q.getpalette()
    if len(raw) < colors * 3:
        raise ValueError(f"quantiser returned {len(raw)//3} entries, wanted {colors}")
    pal = np.array(raw[:colors * 3], dtype=np.uint8).reshape(-1, 3)
    pal = np.vstack([np.zeros((1, 3), np.uint8), pal, np.zeros((1, 3), np.uint8)])
    idx = np.asarray(q, dtype=np.uint8) + np.uint8(1)
    if idx.min() < 1 or idx.max() > colors:
        raise ValueError("index 0 (transparency key) or 255 (black) got used")
    return pal, idx

def palettise(rgb, ref):
    return [np.asarray(Image.fromarray(r).quantize(palette=ref, dither=Image.NONE),
                       dtype=np.uint8) for r in rgb]


def blockify(idx, pal, bw=BW, bh=BH):
    """(H,W) palette indices -> (nblocks, bh*bw*3) float32 luma-weighted RGB."""
    H, W = idx.shape
    rgb = pal[idx].astype(np.float32) * LUMA          # weight once, up front
    b = rgb.reshape(H // bh, bh, W // bw, bw, 3).transpose(0, 2, 1, 3, 4)
    return b.reshape(-1, bh * bw * 3)


def kmeans(X, k, iters=24, seed=0):
    """Chunked Lloyd's algorithm. k-means++ style seeding, deterministic."""
    rng = np.random.default_rng(seed)
    n = X.shape[0]
    if n <= k:
        return X.copy(), np.arange(n)
    # seed: farthest-point sampling on a random subsample (cheap k-means++)
    sub = X[rng.choice(n, min(n, 20000), replace=False)]
    C = np.empty((k, X.shape[1]), dtype=np.float32)
    C[0] = sub[rng.integers(len(sub))]
    d2 = ((sub - C[0]) ** 2).sum(1)
    for i in range(1, k):
        C[i] = sub[np.argmax(d2)]
        d2 = np.minimum(d2, ((sub - C[i]) ** 2).sum(1))
    lab = None
    for _ in range(iters):
        lab = assign(X, C)
        newC = C.copy()
        cnt = np.bincount(lab, minlength=k)
        s = np.zeros_like(C)
        np.add.at(s, lab, X)
        nz = cnt > 0
        newC[nz] = s[nz] / cnt[nz, None]
        # revive dead codewords on the worst-fit blocks
        if (~nz).any():
            err = ((X - newC[lab]) ** 2).sum(1)
            worst = np.argsort(err)[-int((~nz).sum()):]
            newC[~nz] = X[worst]
        if np.allclose(newC, C):
            C = newC; break
        C = newC
    return C, assign(X, C)


# k-means assignment is ~95% of an encode's wall clock (the rest of the encoder,
# rate control included, is about a second for a 120-frame window), so it is
# worth the three details below.  All three are EXACT: the labels are unchanged
# bit for bit, which is what lets the containers this encoder emits stay
# byte-identical across the change.
_ASSIGN_CHUNK   = 2048                     # measured: see the note in assign()
_ASSIGN_THREADS = min(8, (os.cpu_count() or 1))


def _assign_range(X, CT, Cn, out, lo, hi, chunk):
    for i in range(lo, hi, chunk):
        j = min(i + chunk, hi)
        d = Cn[None, :] - 2.0 * (X[i:j] @ CT)      # + |x|^2, constant per row
        out[i:j] = np.argmin(d, axis=1)


def assign(X, C, chunk=_ASSIGN_CHUNK, threads=None):
    """Nearest centroid, chunked to bound memory.

    Three things make this 5.9x faster than the obvious version, measured on
    the 120-frame Singe window (1,474,560 2x2 blocks against k=256), and none
    of them changes a label:

      * `C.T` is a VIEW, and a non-contiguous right-hand operand makes BLAS
        copy it per chunk: 1.83s -> 1.02s just from materialising it once.
      * chunk=2048, not 8192.  The temporary is (chunk, k) float32 and the win
        is cache residency, not memory: 8192 is 1.01s, 32768 is 2.80s.
      * the chunk loop is embarrassingly parallel and numpy releases the GIL in
        both the matmul and the argmin, so a plain thread pool scales it:
        1.02 -> 0.31s on 8 threads.  Partitioning by row cannot change an
        argmin, so the labels are identical to the serial ones -- asserted in
        tools/analysis/09_ratectl_drift.py by the fact that every container
        this encoder emits still hashes the same.
    """
    Cn = (C ** 2).sum(1)
    CT = np.ascontiguousarray(C.T)
    out = np.empty(X.shape[0], dtype=np.int32)
    n = X.shape[0]
    nt = _ASSIGN_THREADS if threads is None else threads
    if nt <= 1 or n < 4 * chunk:
        _assign_range(X, CT, Cn, out, 0, n, chunk)
        return out
    bnd = [(n * i) // nt for i in range(nt + 1)]
    with ThreadPoolExecutor(nt) as ex:
        list(ex.map(lambda ab: _assign_range(X, CT, Cn, out, ab[0], ab[1], chunk),
                    zip(bnd[:-1], bnd[1:])))
    return out


def snap_codebook(C, pal, bw=BW, bh=BH):
    """Centroids (luma-weighted RGB) -> legal palette indices, as the ROM stores them."""
    cb_rgb = C.reshape(-1, bh * bw, 3) / LUMA        # undo the weighting
    palw = pal.astype(np.float32) * LUMA
    flat = (cb_rgb * LUMA).reshape(-1, 3)
    d = (flat ** 2).sum(1)[:, None] - 2 * (flat @ palw.T) + (palw ** 2).sum(1)[None, :]
    return np.argmin(d, axis=1).astype(np.uint8).reshape(-1, bh * bw)


def unblockify(labels, cb_idx, H, W, bw=BW, bh=BH):
    blocks = cb_idx[labels].reshape(H // bh, W // bw, bh, bw)
    return blocks.transpose(0, 2, 1, 3).reshape(H, W)


def psnr(a, b):
    mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
    return 99.0 if mse == 0 else 10 * np.log10(255.0 ** 2 / mse)


def encode_scene(frames_dir, k=256, bw=BW, bh=BH, iters=24):
    rgb = load_frames(frames_dir)
    H, W = rgb[0].shape[:2]
    ref, pal = scene_palette(rgb)
    idx = palettise(rgb, ref)

    X = np.concatenate([blockify(i, pal, bw, bh) for i in idx])
    C, _ = kmeans(X, k, iters)
    cb_idx = snap_codebook(C, pal, bw, bh)

    # re-assign against the SNAPPED codebook: that's what the decoder can produce
    Csnap = (pal[cb_idx].astype(np.float32) * LUMA).reshape(k, -1)
    recon, labels = [], []
    for i in idx:
        lab = assign(blockify(i, pal, bw, bh), Csnap)
        labels.append(lab)
        recon.append(unblockify(lab, cb_idx, H, W, bw, bh))
    return dict(rgb=rgb, pal=pal, idx=idx, cb_idx=cb_idx, labels=labels,
                recon=recon, H=H, W=W, k=k, bw=bw, bh=bh)


def report(r, name=""):
    pal, idx, recon = r["pal"], r["idx"], r["recon"]
    src8 = [pal[i] for i in idx]
    vq8 = [pal[i] for i in recon]
    orig = r["rgb"]
    p_pal = np.mean([psnr(o, s) for o, s in zip(orig, src8)])
    p_vq = np.mean([psnr(o, v) for o, v in zip(orig, vq8)])
    p_vq_only = np.mean([psnr(s, v) for s, v in zip(src8, vq8)])
    nb = (r["H"] // r["bh"]) * (r["W"] // r["bw"])
    idxbits = int(np.ceil(np.log2(r["k"])))
    keyf = nb * idxbits / 8
    # block-delta cost: how many block indices change frame to frame
    ch = [np.count_nonzero(r["labels"][i] != r["labels"][i - 1]) / nb
          for i in range(1, len(r["labels"]))]
    print(f"--- {name} k={r['k']} block={r['bw']}x{r['bh']} ---")
    print(f"  palette-only PSNR : {p_pal:5.2f} dB   (floor: 256c is the best we can do)")
    print(f"  after VQ     PSNR : {p_vq:5.2f} dB   (loss from VQ alone: {p_pal-p_vq:.2f} dB)")
    print(f"  VQ vs palettised  : {p_vq_only:5.2f} dB")
    print(f"  blocks/frame      : {nb}  keyframe {keyf:.0f} B  codebook {r['k']*r['bw']*r['bh']} B")
    if ch:
        print(f"  blocks changed/frm: mean {100*np.mean(ch):5.1f}%  p90 {100*np.percentile(ch,90):5.1f}%")
    return dict(p_pal=p_pal, p_vq=p_vq, nb=nb, keyf=keyf,
                chg=np.mean(ch) if ch else 0, chg90=np.percentile(ch,90) if ch else 0)


def zoom(a, f=3):
    return np.repeat(np.repeat(a, f, axis=0), f, axis=1)


def compare_png(r, frame, out, f=3):
    pal = r["pal"]
    src = pal[r["idx"][frame]]
    vq = pal[r["recon"][frame]]
    orig = r["rgb"][frame]
    gap = np.full((src.shape[0] * f, 4, 3), 40, dtype=np.uint8)
    strip = np.concatenate([zoom(orig, f), gap, zoom(src, f), gap, zoom(vq, f)], axis=1)
    Image.fromarray(strip).save(out)
    return out