#!/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 # --------------------------------------------------------------------------- # CYCLE cost of each mode, per block, MEASURED on the 68000 (FINDINGS 28.2, # tools/bench/decode.lua). This is the other axis: `lam` prices bytes, `mu` # prices cycles, and the two are not proportional -- V4 is 4x a V1 block in # bytes and 1.49x in cycles. # # SKIP IS NOT A CONSTANT, and it is the one trap in here. A SKIP block costs # 13.25 cycles when all four blocks sharing its header byte are SKIP (one # `tst.b` clears the group) and ~45 when it sits in a mixed byte -- so its # price depends on its NEIGHBOURS, which a per-block lagrangian cannot see. # The way out is that the two uses do not need the same number: # * `decide` uses C_SKIP_RANK purely to RANK modes within a block. SKIP is # the cheapest mode either way, so the choice only scales the incentive: # the V1-SKIP gap moves 12% between the two candidates. # * `cycles()` scores a WHOLE frame with the exact clustered rule, and that # is what the rate controller bisects against. Nothing downstream of the # mode decision uses the ranking constant. C_V1, C_V4, C_RAW = 299.9, 448.2, 400.4 C_SKIP_CLUSTERED = 53.0 / 4 # all-SKIP header byte: one tst.b for four # C_SKIP_MIXED WAS THE ONE CONSTANT HERE THAT HAD NEVER BEEN MEASURED. It was # 45.0, hand-derived, from session 7 until session 12 measured it -- and it was # 18% low. Every other figure in this table comes from a synthetic frame of a # single mode, and there was no such frame for a SKIP in a MIXED byte, because # a frame of nothing but mixed SKIPs cannot exist: the byte has to hold a coded # block for the SKIP to be mixed at all. # # tools/bench/prep_dlx.py now emits four that bracket it -- (3 SKIP + 1 V1), # (1 SKIP + 3 V1), (3 SKIP + 1 RAW), (1 SKIP + 3 RAW), each with the header byte # ROTATED through all four positions so no mode is pinned to the free `lsr` # slot -- and each pair solves for the SKIP cost and its partner's together: # # MAME C68K (the partner solves back to its own anchored # V1 pair 55.03 56.50 value to 0.2%, which is what says the pair # RAW pair 55.83 56.50 is measuring the SKIP and not absorbing it) # # 55.0 is taken because every other constant here is MAME's; C68K reads V4 and # RAW 3.2-3.5% higher on pure frames too, which is FINDINGS 37's known table # spread and not a property of mixed bytes. # # It matters more than 10 clocks a block sounds, because a v7 SPAN marks its # run SKIP: a spanned container is made largely of mixed SKIPs, so this is the # dominant population in exactly the frames spans are judged on. FINDINGS 41.5. C_SKIP_MIXED = 55.0 # a SKIP block inside a mixed byte, MEASURED C_SKIP_RANK = C_SKIP_CLUSTERED # ranking only -- see above MODE_CYCLES = np.array([C_SKIP_RANK, C_V1, C_V4, C_RAW], dtype=np.float64) FRAME_CYCLES_12FPS = 10_000_000 / 12.0 # 833,333, x68k.cpp:1133 def cycles(mode): """Exact decode cost of one frame's mode map, in 68000 cycles. Single source of truth: tools/analysis/11_cpu_budget.py imports this, and it reproduces the four frames timed on the 68000 to within 1 point (FINDINGS 28.2). Instruction cycles against zero-wait-state memory, so a LOWER BOUND like every 68000 figure since FINDINGS 24.""" g = np.asarray(mode).reshape(-1, 4) # one header byte = four blocks allskip = (g == 0).all(1) c = allskip.sum() * 4 * C_SKIP_CLUSTERED mm = g[~allskip] c += (mm == 0).sum() * C_SKIP_MIXED c += (mm == 1).sum() * C_V1 c += (mm == 2).sum() * C_V4 c += (mm == 3).sum() * C_RAW return float(c) 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, mu=0.0, byte_clk=0.0): """Lagrangian mode decision at one lam and one mu. Returns (mode, bytes). Minimises `distortion + lam*bytes + mu*(decode cycles + byte_clk*bytes)` per block. `mu=0` is the byte-only decision every session before 8 made. `byte_clk` is what a DELIVERED byte costs the 68000 in clocks of the same budget mu is bisected against -- ratectl.DISK_CLK_BYTE, 5.0 on the single-address row of FINDINGS 43.2. It defaults to 0, which reproduces sessions 8-14 exactly, and it is the last place in the encoder where a byte was still free: FINDINGS 43.6 put the disk debit in the rate controller's FIT TEST, but the decision underneath it still ranked modes as though the 16 bytes of a RAW block cost nothing to deliver. THAT INVERTS FINDINGS 28.8. RAW is 400.4 cycles against V4's 448.2, so with a free byte, raising mu buys cycles by moving V4 -> RAW. Priced, a RAW block costs `400.4 + 16c` and a V4 block `448.2 + 4c`, which cross at **c = 3.98 clocks/byte** -- and 43.1's floor argument (a 68000 bus cycle is four clocks and the SPC hands over one byte per cycle) says c >= 4 on any real machine. So on hardware mu's escape hatch was never there: it was spending 12 clocks of bus to save 47.8 of CPU. Cheap by design: no painting, no image-sized work. A search calls this a dozen times per lam step and paints once.""" ib = ctx["idx_bytes"] s = ctx["sym"] mc = mu * MODE_CYCLES lb = lam + mu * byte_clk # what one payload byte costs, both budgets cost = np.stack([ctx["eS"] + mc[0], s["e1"] + lb * (1.0 * ib) + mc[1], s["e4"] + lb * (4.0 * ib) + mc[2], np.full(ctx["nb"], lb * RAW_BYTES + mc[3])]) 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, mu=0.0): """One frame at one lam against one previous reconstruction.""" ctx = frame_ctx(m, f, prev, idx_bytes) mode, sz = decide(ctx, lam, mu) 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())