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