drwily/Dragon-s-Lair-X68k
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rate control: rebuilt per-frame, wired in, and gated at zero drift

Browse Source 
FINDINGS 26 stopped the session-5 rate controller before it shipped: it built a
lam-ladder of independent whole-sequence encodes and picked frames off it, so
SKIP blocks referenced reconstructions the decoder never saw -- 111 of 120
frames drifted. The fix is the structural one 26.1 said it had to be.

vq_hybrid is now frame-drivable -- frame_ctx / decide / paint -- and encode() is
a thin loop over it. Rate control drives the same three calls, bisects lam per
frame under the leaky bucket, and feeds back the frame it actually emitted. The
desync has no way to occur, and 09_ratectl_drift.py goes 111/120 -> 0/120. That
test is now part of check.sh, which is ~2 min rather than ~40 s.

Both overshoots on the worst sustained window are closed for under 1 dB, totals
including audio: sasi 137.4 -> 109.5 KB/s (-0.60 dB), scsi 381.6 -> 280.0 KB/s
(-0.91 dB). Zero frames hit the lam=800 cliff, so nothing was destroyed to get
there. Rate control also makes the display path cheaper -- scsi's median drops
53.6% -> 47.1% -- because raising lam moves blocks to SKIP and V1.

Two knobs measured rather than guessed. --rc-floor is worth 0.00 dB on that
window and defaults to the profile lam, so rate control cannot regress content
that already fits. --prefill defaults to 0 and is documented as a trap: it buys
a permission to overshoot of exactly bucket/nframes, and on a 14-frame clip it
disables the controller outright.

FINDINGS 26.5 was wrong in both halves and 27.6 records it. _paint was not the
bottleneck (14% of a frame, though vectorising it was still right at 17.1x) and
the ladder was never "minutes" -- those were k-means in build(). What makes
per-frame rate control affordable is that VQ.assign depends on neither lam nor
prev, so it is cached one frame deep: a 12-step search over 120 frames costs
0.31 s against 49.1 s.

Also caught: fixed-lam sasi was already 5% over target on 00020, the clip
everyone called easy. Nothing noticed because the profile table quotes PSNR and
not bitrate.

check.sh: ALL GREEN.

Claude-Session: https://claude.ai/code/session_01194oWYW8DQXK1SZ2DnChW6
This commit is contained in:
prosolis
2026-08-23 14:36:45 -07:00
parent 145753c0bf
commit 497f88b945
9 changed files with 653 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/encoder/encode.py
+56 -9
View File
@@ -3,6 +3,14 @@
python3 tools/encoder/encode.py <frames_dir> <out.dlx> [--profile sasi|scsi]
[--lam N] [--fps 12] [--preview out.png]
[--fixed-lam] [--rc-floor profile|open]
Rate control is ON by default: lam is bisected per frame under a leaky bucket
so the profile's bitrate is a ceiling rather than an average hope. `--fixed-lam`
restores session 5's behaviour, which overshoots by 18-34% on sustained action
(FINDINGS 25.3). `--rc-floor` picks the quality floor: `profile` (default) never
spends more than the fixed-lam profile would, so it can only ever help; `open`
lets quiet frames spend the whole allowance and lands the mean ON target.
Container (little-endian is WRONG here -- the 68000 is big-endian, so every
multi-byte field is big-endian and the decoder can read it with a plain move.w):
@@ -79,6 +87,16 @@ def main():
ap.add_argument("--lam", type=float, default=None)
ap.add_argument("--fps", type=int, default=12)
ap.add_argument("--iters", type=int, default=16)
ap.add_argument("--fixed-lam", action="store_true",
help="disable rate control (session 5 behaviour)")
ap.add_argument("--rc-floor", choices=("profile", "open"), default="profile",
help="quality floor for rate control")
ap.add_argument("--bucket-frames", type=int, default=8,
help="leaky-bucket depth, in frame budgets")
ap.add_argument("--prefill", type=float, default=0.0,
help="how full the player's buffer is assumed to be at "
"scene start, as a fraction of the bucket (0 = cold "
"buffer after a seek, the conservative assumption)")
ap.add_argument("--preview")
a = ap.parse_args()
@@ -87,26 +105,43 @@ def main():
k1, k4 = prof["k1"], prof["k4"]
_IDX_BYTES = 1 if max(k1, k4) <= 256 else 2
# An explicit --lam is a request for that lam, so it implies --fixed-lam.
rc = not (a.fixed_lam or a.lam is not None)
lam_lo = lam if a.rc_floor == "profile" else 1.0
print(f"profile {a.profile}: {prof['desc']}")
print(f" target {prof['kbps']} KB/s, lam={lam}, k1={k1} k4={k4}, "
f"{_IDX_BYTES}-byte indices")
if rc:
print(f" target {prof['kbps']} KB/s CEILING, rate-controlled: "
f"lam bisected per frame in [{lam_lo:g}, {RC.LAM_CLIFF:g}], "
f"{a.bucket_frames}-frame bucket")
else:
print(f" target {prof['kbps']} KB/s, FIXED lam={lam} (no rate control)")
print(f" k1={k1} k4={k4}, {_IDX_BYTES}-byte indices")
m = H.build(a.frames_dir, k1=k1, k4=k4, iters=a.iters)
enc = H.encode(m, lam=lam)
if rc:
enc = RC.encode_rate_controlled(m, prof["kbps"], fps=a.fps,
bucket_frames=a.bucket_frames,
lam_lo=lam_lo, prefill=a.prefill)
else:
enc = H.encode(m, lam=lam)
r = H.evaluate(m, enc, fps=a.fps)
H_, W_ = m["H"], m["W"]; nbx = W_ // 4
pal, idx = m["pal"], m["idx"]
# re-derive the per-frame symbols the same way encode() did
# The encoder hands back the symbols it actually chose. Re-deriving them
# here (as session 5 did) is a second chance to disagree with the encoder,
# and with per-frame rate control the mode map is no longer reproducible
# from a single lam anyway.
frames = []
for f, im in enumerate(idx):
B1 = H.blocks_of(im, pal, 4, 4); l1 = VQ.assign(B1, m["C1s"])
B4 = H.blocks_of(im, pal, 2, 2); l4 = VQ.assign(B4, m["C4s"])
q = H._group_2x2_into_4x4(np.arange(len(l4)), W_)
l4g = l4[q].reshape(-1, 4)
mode = enc["modes"][f]
frames.append(pack_modes(mode) + frame_payload(mode, l1, l4g, im, nbx))
frames.append(pack_modes(mode)
+ frame_payload(mode, enc["l1"][f], enc["l4g"][f], im, nbx))
# the rate controller budgets exactly these bytes -- if that ever drifts
# from the container, every bitrate figure below is fiction
assert len(frames[-1]) == enc["sizes"][f], (f, len(frames[-1]), enc["sizes"][f])
palette = m["pal"][:256]
if len(palette) < 256:
@@ -138,6 +173,18 @@ def main():
f"loss {r['loss']:.2f} dB")
print(f" modes: SKIP {r['skip']:.1f}% V1 {r['v1']:.1f}% "
f"V4 {r['v4']:.1f}% RAW {r['raw']:.1f}%")
if rc:
rr = RC.summarise(m, enc, prof["kbps"], fps=a.fps)
lm = enc["lam"]
print(f" rate control: per-frame budget {enc['budget']:.0f} B, "
f"bucket {enc['cap']:.0f} B ({a.bucket_frames} frames), "
f"prefill {100*a.prefill:.0f}%")
print(f" lam: min {lm.min():.1f} median {rr['lam_med']:.1f} "
f"p90 {rr['lam_p90']:.1f} max {rr['lam_max']:.1f}")
print(f" frames over the per-frame budget (banked by the bucket): "
f"{rr['over']:.0f}%")
print(f" frames that could not fit even at the lam={RC.LAM_CLIFF:g} "
f"cliff: {rr['overrun']}/{len(lm)}")
# PER-FRAME non-SKIP distribution. The mean above cannot answer the
# decoder-architecture question (FINDINGS 24.5): decode-direct-to-GVRAM
tools/encoder/ratectl.py
+118 -44
View File
@@ -12,11 +12,12 @@ without that, quiet frames waste budget and action frames stay ugly.
The ceiling is HARD: the 68000 streams at a fixed rate off the disk, and a frame
that overruns is a dropped frame, not a slow frame.
STATUS, session 5: this module is written but STILL NOT WIRED INTO encode.py,
and FINDINGS 25.3 measured both profiles overshooting their targets by 18% and
34% on the worst sustained window because of that. Before wiring it up, read
the correctness note on encode_rate_controlled() -- the lam-ladder approach it
uses is not sound against a temporally recursive encoder.
STATUS, session 6: WIRED IN and sound. `encode.py` rate-controls by default
for a profile; `--fixed-lam` restores the old behaviour. The lam-ladder of
session 5 was replaced by a per-frame bisection that drives the encoder one
frame at a time and feeds back the frame it actually emitted -- see
encode_rate_controlled(), and FINDINGS 26 for why the ladder could not be
fixed by tuning. Regression test: tools/analysis/09_ratectl_drift.py.
"""
import numpy as np
import vq_hybrid as H
@@ -49,16 +50,22 @@ import vq_hybrid as H
# The rates below are therefore RAW payload, no entropy coding.
#
# The two profiles are the SAME codec, decoder and bitstream -- only `lam` differs.
# `lam` here is a FLOOR, not a setting: encode.py rate-controls by default and
# bisects lam per frame in [lam, LAM_CLIFF] to keep under `kbps`. The floor is
# what a quiet frame is allowed to spend, so rate control can only ever spend
# less than session 5's fixed-lam encoder did. FINDINGS 27.
PROFILES = {
"sasi": dict(kbps=110, lam=60.0, k1=256, k4=256,
desc="stock 10MHz ACE/EXPERT, SASI",
quality="36.9 dB on 00020 / 29.6 dB on 00146 / 27.8 dB on the "
"Singe window, where it overshoots to 129.6 KB/s",
quality="36.9 dB on 00020 / 29.6 dB on 00146 / 27.2 dB on the "
"Singe window at 109.5 KB/s (session 5's fixed lam "
"gave 27.8 dB there, but at 137.4 KB/s)",
util="~105 KB/s = 35% of the pessimistic 300 KB/s SASI figure"),
"scsi": dict(kbps=280, lam=10.0, k1=256, k4=256,
desc="Super/XVI, or CZ-6BS1 board in a 10MHz machine",
quality="39.4 dB on 00020 / 32.3 dB on 00146 / 30.8 dB on the "
"Singe window, where it overshoots to 373.8 KB/s",
quality="39.4 dB on 00020 / 32.3 dB on 00146 / 29.9 dB on the "
"Singe window at 280.0 KB/s (session 5's fixed lam "
"gave 30.8 dB there, but at 381.6 KB/s)",
util="~275 KB/s = 28% of the 1 MB/s SCSI folklore figure"),
}
# lam=0 is PIXEL-EXACT against the palettised frame (0.00 dB loss) at ~450 KB/s
@@ -68,6 +75,12 @@ PROFILES = {
# lam=0 and the port ships transparent video. That decision is waiting on a
# measurement, not on a design choice.
# Hard ceiling on the rate-control search. FINDINGS 15 puts the quality cliff
# between lam=800 and lam=2000. Above it a frame has not been rate-controlled,
# it has been destroyed, so the search stops here and lets the frame overrun
# instead (FINDINGS 26.2). The old ladder ran to lam=2e5, 250x past shippable.
LAM_CLIFF = 800.0
AUDIO_KBPS = 7.8 # MSM6258 ADPCM 15.6kHz mono -- comes out of the same budget
@@ -76,39 +89,92 @@ def frame_budget(kbps, fps=12, audio=AUDIO_KBPS):
return (kbps - audio) * 1024.0 / fps
def _search_lam(ctx, allow, lam_lo, lam_hi, iters=12):
"""Smallest lam (=> best quality) whose frame fits `allow` bytes.
Payload size is non-increasing in lam -- raising lam can only move a block
to a mode that costs no more -- so bisection is sound. Geometric bisection,
because lam spans three decades and the interesting range is multiplicative.
Returns (lam, mode, size, overrun). `overrun` is True when even lam_hi does
not fit: that frame is emitted over budget on purpose. Past the FINDINGS 15
cliff a frame is not rate-controlled, it is destroyed, so a visible overrun
is the better failure (FINDINGS 26.2)."""
mode, sz = H.decide(ctx, lam_lo)
if sz <= allow:
return lam_lo, mode, sz, False
mode_hi, sz_hi = H.decide(ctx, lam_hi)
if sz_hi > allow:
return lam_hi, mode_hi, sz_hi, True
lo, hi = lam_lo, lam_hi # lo does not fit, hi does
best = (lam_hi, mode_hi, sz_hi)
for _ in range(iters):
mid = float(np.sqrt(lo * hi))
mode_m, sz_m = H.decide(ctx, mid)
if sz_m <= allow:
hi = mid; best = (mid, mode_m, sz_m)
else:
lo = mid
return best[0], best[1], best[2], False
def encode_rate_controlled(m, target_kbps, fps=12, bucket_frames=8,
lam_lo=1.0, lam_hi=2e5, steps=9, verbose=False):
lam_lo=1.0, lam_hi=LAM_CLIFF, prefill=0.0,
steps=None, verbose=False):
"""Per-frame lam search under a leaky bucket, driving the encoder ONE FRAME
AT A TIME and feeding back the frame actually emitted.
That feedback is the whole point. The previous implementation encoded the
sequence once per lam and then picked frames off the resulting ladder; the
codec is temporally recursive, so frames picked from different rungs
reference reconstructions the decoder never saw -- 111 of 120 frames drifted,
worst frame 43.4% (FINDINGS 26.1). `tools/analysis/09_ratectl_drift.py` is
the regression test and must report zero drifting frames.
lam_lo is a QUALITY FLOOR, not a starting guess: rate control here only ever
spends less than the fixed-lam profile, never more, so it cannot regress
content that already fits. Pass lam_lo=1.0 to let quiet frames spend the
whole allowance instead.
`prefill` is how full the player's buffer is assumed to be when the scene
starts, as a fraction of the bucket. 0.0 (the default) is the conservative
assumption -- a cold buffer after a seek -- and is what FINDINGS 21 verified
needs no prefill to avoid underflow. It costs a startup transient: the first
`bucket_frames` frames cannot draw on a bank they have not accumulated yet,
so a clip shorter than a few bucket depths lands UNDER target. That is an
artefact of the clip length, not of the content; see FINDINGS 27.5.
DO NOT raise `prefill` to make a target look met. It works by permitting an
overshoot of cap/nframes: measured, prefill=1.0 takes the Singe window from
109.5 to 116.3 KB/s against a 110 ceiling, and on a 14-frame clip it
disables rate control entirely because the bucket is larger than the clip.
`steps` is accepted and ignored -- there is no ladder any more.
"""
if steps is not None and verbose:
print(" note: `steps` is ignored; lam is now bisected per frame")
budget = frame_budget(target_kbps, fps)
bucket = 0.0 # banked bytes, capped at bucket_frames*budget
cap = bucket_frames * budget
out_recon, out_modes, out_sizes, out_lam = [], [], [], []
# encode() is whole-sequence; drive it per-lam and pick per frame.
# Cheaper than re-running the whole encoder per frame: precompute the ladder.
ladder = []
lams = np.geomspace(lam_lo, lam_hi, steps)
for lam in lams:
e = H.encode(m, lam=float(lam))
ladder.append(e)
if verbose:
print(f" lam={lam:9.0f} mean {e['sizes'].mean():6.0f} B/frame")
nf = len(m["idx"])
for f in range(nf):
bucket = prefill * cap # banked bytes; bounded by the player's buffer both ways
out = dict(recon=[], modes=[], sizes=[], lam=[], l1=[], l4g=[], overrun=[])
prev = None
for f in range(len(m["idx"])):
ctx = H.frame_ctx(m, f, prev)
allow = budget + bucket
# cheapest lam (highest quality) whose size fits the allowance
pick = len(lams) - 1
for i in range(len(lams)):
if ladder[i]["sizes"][f] <= allow:
pick = i; break
sz = ladder[pick]["sizes"][f]
bucket = min(cap, bucket + budget - sz)
out_recon.append(ladder[pick]["recon"][f])
out_modes.append(ladder[pick]["modes"][f])
out_sizes.append(sz); out_lam.append(lams[pick])
return dict(recon=out_recon, modes=out_modes, sizes=np.array(out_sizes),
lam=np.array(out_lam), nb=ladder[0]["nb"], budget=budget)
lam, mode, sz, ovr = _search_lam(ctx, allow, lam_lo, lam_hi)
rec = H.paint(m, ctx, mode)
bucket = float(np.clip(bucket + budget - sz, -cap, cap))
out["recon"].append(rec); out["modes"].append(mode)
out["sizes"].append(sz); out["lam"].append(lam); out["overrun"].append(ovr)
out["l1"].append(ctx["sym"]["l1"]); out["l4g"].append(ctx["sym"]["l4g"])
prev = rec
if verbose:
print(f" f{f:04d} lam={lam:8.2f} {sz:7.0f} B "
f"(allow {allow:7.0f}){' OVER' if ovr else ''}")
return dict(recon=out["recon"], modes=out["modes"],
sizes=np.array(out["sizes"]), lam=np.array(out["lam"]),
l1=out["l1"], l4g=out["l4g"], overrun=np.array(out["overrun"]),
nb=m["nb"], budget=budget, cap=cap)
def summarise(m, enc, target_kbps, fps=12):
@@ -120,9 +186,17 @@ def summarise(m, enc, target_kbps, fps=12):
pp = np.mean([VQ.psnr(o, v) for o, v in zip(m["rgb"], src)])
sz = enc["sizes"]
mo = np.concatenate(enc["modes"])
return dict(target=target_kbps, psnr=p, pal=pp, loss=pp - p,
mean_B=sz.mean(), max_B=sz.max(), budget=enc["budget"],
kbps=sz.mean() * fps / 1024 + AUDIO_KBPS,
over=100.0 * np.mean(sz > enc["budget"]),
skip=100 * (mo == 0).mean(), v1=100 * (mo == 1).mean(),
v4=100 * (mo == 2).mean())
d = dict(target=target_kbps, psnr=p, pal=pp, loss=pp - p,
mean_B=sz.mean(), max_B=sz.max(), budget=enc.get("budget", 0.0),
kbps=sz.mean() * fps / 1024 + AUDIO_KBPS,
over=100.0 * np.mean(sz > enc.get("budget", np.inf)),
skip=100 * (mo == 0).mean(), v1=100 * (mo == 1).mean(),
v4=100 * (mo == 2).mean(), raw=100 * (mo == 3).mean())
if "lam" in enc:
lam = enc["lam"]
d.update(lam_med=float(np.median(lam)), lam_max=float(lam.max()),
lam_p90=float(np.percentile(lam, 90)),
# a frame that could not fit even at the cliff: emitted over
# budget on purpose rather than destroyed
overrun=int(np.asarray(enc.get("overrun", [])).sum()))
return d
tools/encoder/vq_hybrid.py
+164 -68
View File
@@ -17,6 +17,19 @@ 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
@@ -24,6 +37,11 @@ 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
def blocks_of(idx, pal, bw, bh):
return VQ.blockify(idx, pal, bw, bh)
@@ -46,68 +64,163 @@ def build(frames_dir, k1=256, k4=256, iters=16, lam=0.0):
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)
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):
"""Lagrangian mode decision at one lam. Returns (mode, payload bytes).
Cheap by design: no painting, no image-sized work. A lam search calls this
a dozen times per frame and paints once."""
ib = ctx["idx_bytes"]
s = ctx["sym"]
cost = np.stack([ctx["eS"],
s["e1"] + lam * (1.0 * ib),
s["e4"] + lam * (4.0 * ib),
np.full(ctx["nb"], lam * RAW_BYTES)])
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):
"""One frame at one lam against one previous reconstruction."""
ctx = frame_ctx(m, f, prev, idx_bytes)
mode, sz = decide(ctx, lam)
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):
"""lam = lagrangian rate weight (bytes -> squared-error units).
"""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.
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."""
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 = 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 = [], [], []
idx_bytes = default_idx_bytes(m)
recon, modes, sizes, l1s, l4gs = [], [], [], [], []
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)
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):
@@ -119,23 +232,6 @@ def _group_2x2_into_4x4(a, W):
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"]]