File size: 3,415 Bytes
0692312 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | #!/usr/bin/env python
"""One grid video over the lambda experiments, each panel captioned with what it is.
Composed here rather than with an ffmpeg filtergraph because the caption is the whole point -- six
unlabelled panels are unreadable, and `drawtext` needs a freetype-enabled build that the bundled
imageio-ffmpeg binary does not reliably have. PIL draws the label, imageio writes the file.
Panels are scaled to a common height and letterboxed to a common width, so a canvas difference
between runs shows as padding instead of silently stretching one panel against the others.
"""
from __future__ import annotations
import argparse
from pathlib import Path
import av
import numpy as np
from PIL import Image, ImageDraw
PANELS = [
("0_bf16", "BF16 reference"),
("1_plain_w4a4", "1. plain W4A4 (lambda=1, no LoRA)"),
("2_lambda_none", "2. lambda=1 + rank-32 LoRA"),
("3_lambda_all", "3. lambda from all tokens"),
("4_lambda_video", "4. lambda from video tokens"),
("5_lambda_text", "5. lambda from text tokens"),
]
def read(path: Path, height: int) -> list[np.ndarray]:
c = av.open(str(path))
v = c.streams.video[0]
v.thread_type, v.thread_count = "NONE", 1
out = []
for f in c.decode(video=0):
a = f.to_ndarray(format="rgb24")
w = int(round(a.shape[1] * height / a.shape[0] / 2)) * 2
out.append(np.asarray(Image.fromarray(a).resize((w, height), Image.BICUBIC)))
c.close()
return out
def label(frame: np.ndarray, text: str) -> np.ndarray:
im = Image.fromarray(frame)
d = ImageDraw.Draw(im, "RGBA")
d.rectangle([0, 0, im.width, 20], fill=(0, 0, 0, 170))
d.text((6, 5), text, fill=(255, 255, 255))
return np.asarray(im)
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--dir", default="out/lambda_exps")
ap.add_argument("--out", default="out/lambda_exps/grid.mp4")
ap.add_argument("--height", type=int, default=300)
ap.add_argument("--fps", type=int, default=24)
args = ap.parse_args()
d = Path(args.dir)
have = [(s, t) for s, t in PANELS if (d / f"{s}.mp4").exists()]
if len(have) < 2:
print(f"only {len(have)} panels present, nothing to stack")
return 0
clips = [read(d / f"{s}.mp4", args.height) for s, _ in have]
n = min(len(c) for c in clips)
w = max(c[0].shape[1] for c in clips)
cols = (len(have) + 1) // 2
import imageio.v2 as imageio
wr = imageio.get_writer(args.out, fps=args.fps, codec="libx264", quality=8,
macro_block_size=2)
for i in range(n):
tiles = []
for (s, t), c in zip(have, clips):
f = label(c[i], t)
if f.shape[1] < w: # letterbox, never stretch
pad = np.zeros((f.shape[0], w, 3), np.uint8)
o = (w - f.shape[1]) // 2
pad[:, o:o + f.shape[1]] = f
f = pad
tiles.append(f)
while len(tiles) % cols:
tiles.append(np.zeros_like(tiles[0]))
rows = [np.concatenate(tiles[k:k + cols], axis=1) for k in range(0, len(tiles), cols)]
wr.append_data(np.concatenate(rows, axis=0))
wr.close()
print(f"wrote {args.out}: {len(have)} panels, {n} frames, {cols} cols")
for s, t in have:
print(f" {s:<16} {t}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|