e4c7287acb
D2c M4, gate 7 (delete). - int_hyper_synth::decode_conditional, the unchecked conditional decode, is deleted; its callers decode through decode_conditional_checked (the lowering, with its check ranges): the int_hyper_decode_frame example (step 3 of tools/wai_sota_vectors.sh), wai_decode_bench's tiny-vector case, int_payload's JSON-equivalence unit tests, tests/int_hyper_synth.rs, int_kernels_vectors.rs, int_video_keyframe.rs, int_video_sota_keyframe.rs, and byte-exact's neural_keyframe_golden (the trap: it gains the check ranges, outputs equal; gate 2's duplicate test through ops::decode goes). - The pretrained synthesis-only and hyper-decoder-only goldens (tests/int_hyper_synth.rs) are also held through the lowering. - The conditional codec's direct guard (legacy's run_checked) is test-only since gate 6. run_net_checked, check_layer and igdn_check stay: their one production caller left is int_mlicpp (M5); run_net and decode stay for the world renderer (M6), and int_hyper_synth's kernel_equivalence with them, since their conv dispatch is still run_net's. Docs say so. - Doc references (int_video, int_video_golden, two tools) and CHANGELOG. Co-Authored-By: Claude Opus 5.5 <[email protected]>
110 lines
4.6 KiB
Python
110 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Encode the NEURAL-KEYFRAME inter-frame video: the I-frame is decoded by the
|
|
integer learned codec `int_hyper`, and the P-frames are integer block-MC residuals
|
|
on top — the whole decode integer, byte-exact on every machine.
|
|
|
|
The I-frame here is the committed tiny `int_hyper` vector (so the conformance is
|
|
self-contained and never-skips); the P-frames translate it and code the residual.
|
|
The Rust decoder composes `int_hyper_synth::decode_conditional_checked` (frame 0) with
|
|
`int_video::decode_pframes` (the rest) — this script emits only the P-frame
|
|
bitstream + the expected per-frame / clip hashes (frame 0's hash is int_hyper's).
|
|
|
|
python tools/wai_int_video_keyframe_encode.py <int_hyper-vectors-dir> <out-dir>
|
|
"""
|
|
import sys, os, json, base64, tempfile
|
|
import numpy as np
|
|
from compressai.ans import RansEncoder
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from wai_int_hyper_tiny_vector import main as gen_tiny
|
|
|
|
BLOCK, SEARCH, Q, NF = 8, 4, 6, 5
|
|
|
|
|
|
def fnv1a(buf):
|
|
h = 14695981039346656037
|
|
for b in buf:
|
|
h = ((h ^ b) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
|
|
return h
|
|
|
|
|
|
def laplacian_cdf(rmax):
|
|
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, h, w):
|
|
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)]
|
|
for y in range(by, min(by + BLOCK, h)):
|
|
sy = min(max(y + dy, 0), h - 1)
|
|
for x in range(bx, min(bx + BLOCK, w)):
|
|
pred[y, x] = ref[sy, min(max(x + dx, 0), w - 1)]
|
|
return pred
|
|
|
|
|
|
def search_mvs(cur, ref, h, w):
|
|
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(ih_dir, outdir):
|
|
os.makedirs(outdir, exist_ok=True)
|
|
# frame 0 = the int_hyper I-frame (regenerate the tiny vector in a temp dir to
|
|
# read its decoded pixels; the committed int_hyper vector itself is untouched).
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
rgb0, H, W = gen_tiny(tmp)
|
|
frame0 = rgb0.astype(np.int64) # HWC int, 0..255
|
|
rmax = (255 + Q - 1) // Q
|
|
cdf, offset = laplacian_cdf(rmax), -rmax
|
|
|
|
frames = [frame0] + [np.roll(frame0, (f, f), axis=(0, 1)) for f in range(1, NF)] # translate
|
|
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], H, W)
|
|
pred = np.stack([mc_predict(recon[-1][..., c], mvs, H, W) 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'{outdir}/video_nk_bitstream.json', 'w'))
|
|
json.dump({'n_frames': NF, 'h': H, 'w': W, 'frame_hashes': hashes, 'clip_hash': clip_hash},
|
|
open(f'{outdir}/video_nk.json', 'w'))
|
|
print(f'int_hyper-keyframe video: {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], sys.argv[2])
|