#!/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