935a83f3bc
Makes wai.video.int_hyper a live, verifiable artifact: a clip whose keyframe is the learned integer codec int_hyper, decoded byte-exact in the browser, with integer motion-compensated P-frames on top — the whole decode integer, every frame hash re-checked against the native reference. - src/video.rs: decodeIntVideoPframes(i_frame, w, h, bitstream) — the P-frame loop seeded with a given I-frame (factored a flatten() helper shared with decodeIntVideo). The composition is decodeIntHyperBin (learned keyframe) ∘ decodeIntVideoPframes. - tools/wai_int_video_sota_keyframe_encode.py: now infers the frame size from frame0.rgb and takes an output prefix, so one tool makes both the 256² conformance vector and the 128² demo clip. - Drives the new "neural keyframe (int_hyper)" mode on wai.transaction.science/video: fetches the served int_hyper 128² model + bitstream + the P-frame stream, decodes the learned keyframe (~2.2 s) + integer P-frames, verifies all 8 frames byte-for-byte and plays the int_hyper photo panning. The default light int_motion clip stays. Co-Authored-By: Claude Opus 4.8 <[email protected]>
110 lines
4.7 KiB
Python
110 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Encode a REAL 256x256 int_hyper-keyframe clip: the I-frame is the SOTA learned
|
|
integer codec int_hyper (its byte-exact 256x256 reconstruction, dumped by the
|
|
`int_hyper_decode_frame` Rust example), and the P-frames pan it with integer block
|
|
motion compensation + integer-quantized residual. Skip-if-absent: needs the SOTA
|
|
int_hyper vector dir (default /tmp/sota_int) holding frame0.rgb.
|
|
|
|
cargo run --release --example int_hyper_decode_frame --features neural_int -- <dir> <dir>/frame0.rgb
|
|
python tools/wai_int_video_sota_keyframe_encode.py <dir>
|
|
|
|
Emits <dir>/video_sota_bitstream.json + <dir>/video_sota.json (the P-frame stream +
|
|
expected per-frame / clip hashes; frame 0's hash is int_hyper's). The Rust test
|
|
`int_video_sota_keyframe` composes int_hyper::decode_conditional with
|
|
int_video::decode_pframes and checks these, skipping when the dir is absent.
|
|
"""
|
|
import sys, os, json, base64
|
|
import numpy as np
|
|
from compressai.ans import RansEncoder
|
|
|
|
H = W = 256 # default; main() infers the real size from frame0.rgb
|
|
BLOCK, SEARCH, Q, NF = 16, 4, 6, 8
|
|
RMAX = (255 + Q - 1) // Q
|
|
|
|
|
|
def fnv1a(buf):
|
|
h = 14695981039346656037
|
|
for b in buf:
|
|
h = ((h ^ b) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
|
|
return h
|
|
|
|
|
|
def laplacian_cdf():
|
|
n = 2 * RMAX + 1
|
|
pmf = np.exp(-np.abs(np.arange(n) - RMAX) / 6.0)
|
|
total = 1 << 16
|
|
f = np.maximum((pmf / pmf.sum() * (total - 1)).round().astype(np.int64), 1)
|
|
f[f.argmax()] += (total - 1) - int(f.sum())
|
|
cdf = [0]
|
|
for v in f:
|
|
cdf.append(cdf[-1] + int(v))
|
|
cdf.append(total)
|
|
return cdf
|
|
|
|
|
|
def mc_predict(ref, mvs):
|
|
pred = np.empty_like(ref)
|
|
nbx = (W + BLOCK - 1) // BLOCK
|
|
for by in range(0, H, BLOCK):
|
|
for bx in range(0, W, BLOCK):
|
|
dy, dx = mvs[(by // BLOCK) * nbx + (bx // BLOCK)]
|
|
ys = np.clip(np.arange(by, min(by + BLOCK, H)) + dy, 0, H - 1)
|
|
xs = np.clip(np.arange(bx, min(bx + BLOCK, W)) + dx, 0, W - 1)
|
|
pred[by:by + len(ys), bx:bx + len(xs)] = ref[np.ix_(ys, xs)]
|
|
return pred
|
|
|
|
|
|
def search_mvs(cur, ref):
|
|
mvs = []
|
|
for by in range(0, H, BLOCK):
|
|
for bx in range(0, W, BLOCK):
|
|
cb = cur[by:by + BLOCK, bx:bx + BLOCK]
|
|
best, bmv = None, (0, 0)
|
|
for dy in range(-SEARCH, SEARCH + 1):
|
|
for dx in range(-SEARCH, SEARCH + 1):
|
|
ys = np.clip(np.arange(by, by + cb.shape[0]) + dy, 0, H - 1)
|
|
xs = np.clip(np.arange(bx, bx + cb.shape[1]) + dx, 0, W - 1)
|
|
sad = int(np.abs(cb - ref[np.ix_(ys, xs)]).sum())
|
|
if best is None or sad < best:
|
|
best, bmv = sad, (dy, dx)
|
|
mvs.append([int(bmv[0]), int(bmv[1])])
|
|
return mvs
|
|
|
|
|
|
def main(d, prefix='video_sota'):
|
|
global H, W
|
|
fp = f'{d}/frame0.rgb'
|
|
if not os.path.exists(fp):
|
|
print(f'skip: {fp} absent (run the int_hyper_decode_frame example first)')
|
|
return
|
|
raw = open(fp, 'rb').read()
|
|
H = W = int(round((len(raw) / 3) ** 0.5)) # square frame inferred from the dump
|
|
frame0 = np.frombuffer(raw, np.uint8).reshape(H, W, 3).astype(np.int64)
|
|
cdf, offset = laplacian_cdf(), -RMAX
|
|
frames = [frame0] + [np.roll(frame0, (3 * f, 2 * f), axis=(0, 1)) for f in range(1, NF)] # pan
|
|
recon, mvs_all, syms, hashes = [frame0], [[]], [], []
|
|
hashes.append(format(fnv1a(np.ascontiguousarray(frame0.astype(np.uint8)).tobytes()), '016x'))
|
|
for f in range(1, NF):
|
|
mvs = search_mvs(frames[f], recon[-1])
|
|
pred = np.stack([mc_predict(recon[-1][..., c], mvs) for c in range(3)], -1)
|
|
rq = np.round((frames[f] - pred) / Q).astype(np.int64).clip(-RMAX, RMAX)
|
|
syms.extend(rq.reshape(-1).tolist())
|
|
rec = np.clip(pred + rq * Q, 0, 255).astype(np.int64)
|
|
recon.append(rec)
|
|
mvs_all.append(mvs)
|
|
hashes.append(format(fnv1a(np.ascontiguousarray(rec.astype(np.uint8)).tobytes()), '016x'))
|
|
|
|
bits = RansEncoder().encode_with_indexes(syms, [0] * len(syms), [cdf], [len(cdf)], [offset])
|
|
clip_hash = format(fnv1a(b''.join(bytes.fromhex(h) for h in hashes)), '016x')
|
|
json.dump({'block': BLOCK, 'q': Q, 'rmax': RMAX, 'offset': offset, 'cdf': cdf,
|
|
'n_frames': NF, 'h': H, 'w': W, 'mvs': mvs_all,
|
|
'rans_b64': base64.b64encode(bytes(bits)).decode()},
|
|
open(f'{d}/{prefix}_bitstream.json', 'w'))
|
|
json.dump({'n_frames': NF, 'h': H, 'w': W, 'frame_hashes': hashes, 'clip_hash': clip_hash},
|
|
open(f'{d}/{prefix}.json', 'w'))
|
|
print(f'int_hyper-keyframe video [{prefix}]: {NF}x {H}x{W} (I-frame={hashes[0]}), {len(bits)} P rANS bytes '
|
|
f'-> clip_hash {clip_hash}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main(sys.argv[1] if len(sys.argv) > 1 else '/tmp/sota_int', sys.argv[2] if len(sys.argv) > 2 else 'video_sota')
|