vt1200/harness/badapple/clip_from_video.py

#!/usr/bin/env python3
"""Turn a video into the raw 1-bpp frames ``make_tiles.py --clip`` takes.

Each output file is one frame: whole 160-byte rows, whole 32-row tile bands,
top-aligned, **LSB first** -- bit 0 of a byte is the row's leftmost pixel, the
order the terminal's scanout reads (``dwt/screenshot.py`` decodes live VRAM with
the same rule).  An MSB-first packer is self-consistent with an MSB-first
checker and mirrors every eight pixels on a display, which is why the order is
spelled out here rather than left to ffmpeg.

The screen is 1280 px wide and the tile grid is 32 rows tall, so a 4:3 source
scaled to 1280x960 fills 30 of them and the last two rows stay black.  Pass
``--height 1024`` to fill the screen and accept 4:3 being stretched.

Requires ``ffmpeg`` on PATH and numpy.

    clip_from_video.py bad-apple.mp4 frames/ --fps 15
    clip_from_video.py bad-apple.mp4 frames/ --fps 12 --threshold 100 --invert
"""
from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys

ROW_PIXELS = 1280
ROW_BYTES = ROW_PIXELS // 8          # 160


def pack_row(pixels) -> bytes:
    """Pack one row of 0/1 pixels, LSB first (pixel 0 -> bit 0 of byte 0)."""
    if not pixels or len(pixels) % 8:
        raise ValueError("a row is a positive multiple of 8 pixels, got %d"
                         % len(pixels))
    out = bytearray(len(pixels) // 8)
    for i, value in enumerate(pixels):
        if value:
            out[i >> 3] |= 1 << (i & 7)
    return bytes(out)


def frames(video, out_dir, fps=15, width=ROW_PIXELS, height=960, threshold=127,
           invert=False, block=1):
    """Yield ``(frame_no, packed_row_bytes)`` for every video frame."""
    import numpy as np

    vf = ("fps=%s,scale=%d:%d:flags=lanczos,format=gray"
          % (fps, width, height))
    if invert:
        vf += ",negate"
    cmd = ["ffmpeg", "-v", "error", "-i", video, "-vf", vf,
           "-f", "rawvideo", "-"]
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, bufsize=10 ** 8)
    size = width * height
    n = 0
    try:
        while True:
            buf = proc.stdout.read(size)
            if len(buf) < size:
                break
            gray = np.frombuffer(buf, dtype=np.uint8).reshape(height, width)
            if block > 1:
                # Stylise: average each block, threshold the average, then draw
                # the block flat.  Fewer, larger shapes mean fewer changed tiles
                # per frame -- measured on Bad Apple at 60 fps: 80 changed tiles
                # unstylised, 54 at 8x8, 29 at 16x16, against a budget of ~44.
                h, w = height // block, width // block
                small = gray.reshape(h, block, w, block).mean(axis=(1, 3))
                bits = np.repeat(np.repeat(small > threshold, block, axis=0),
                                 block, axis=1)
            else:
                bits = gray > threshold
            yield n, np.packbits(bits, axis=1, bitorder="little")
            n += 1
    finally:
        proc.stdout.close()
        proc.wait()


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("video", help="source video (anything ffmpeg reads)")
    ap.add_argument("out", help="directory for the *.bin frames")
    ap.add_argument("--fps", type=float, default=15.0,
                    help="frames per second to extract (default %(default)s)")
    ap.add_argument("--width", type=int, default=ROW_PIXELS,
                    help="pixels per row (default %(default)s)")
    ap.add_argument("--height", type=int, default=960,
                    help="rows to extract: 960 fills 30 of the 32 tile bands, "
                         "1024 fills the screen (default %(default)s)")
    ap.add_argument("--threshold", type=int, default=127,
                    help="grey level above which a pixel is ink "
                         "(default %(default)s)")
    ap.add_argument("--block", type=int, default=1,
                    help="stylise to NxN blocks (default 1 = no stylisation)")
    ap.add_argument("--invert", action="store_true",
                    help="treat bright pixels as paper (negate first)")
    args = ap.parse_args(argv)

    if shutil.which("ffmpeg") is None:
        print("error: ffmpeg is not on PATH", file=sys.stderr)
        return 2
    if args.width % 8 or args.width > ROW_PIXELS or args.width < 8:
        print("error: --width must be a multiple of 8, at most %d" % ROW_PIXELS,
              file=sys.stderr)
        return 2
    if args.block < 1 or args.block > 64:
        print("error: --block must be 1..64", file=sys.stderr)
        return 2
    if args.height % 32 or args.height > 1024:
        print("error: --height must be a multiple of 32, at most 1024",
              file=sys.stderr)
        return 2

    os.makedirs(args.out, exist_ok=True)
    for name in os.listdir(args.out):
        if name.endswith(".bin"):
            os.remove(os.path.join(args.out, name))

    written = 0
    for n, packed in frames(args.video, args.out, fps=args.fps,
                            width=args.width, height=args.height,
                            threshold=args.threshold, invert=args.invert,
                            block=args.block):
        path = os.path.join(args.out, "%05d.bin" % n)
        with open(path, "wb") as handle:
            handle.write(packed.tobytes())
        written += 1
        if written % 100 == 0:
            print("  %d frames" % written, file=sys.stderr)
    print("%d frames of %dx%d at %s fps -> %s"
          % (written, args.width, args.height, args.fps, args.out))
    return 0


if __name__ == "__main__":
    sys.exit(main())