noaa-goes-visualization/make_comparison_video.py

#!/usr/bin/env python
"""Stack the pristine, current-pipeline and new-pipeline streams side by side.

`bench.py render` writes one composite per timestamp per variant.  This turns three
such streams into a single video with the panes labelled, so the repair can be judged
against both the truth and what the pipeline does today:

    bench.py render --case C --variant pristine --out /tmp/streams
    bench.py render --case C --variant today    --out /tmp/streams
    bench.py render --case C --variant new      --out /tmp/streams
    make_comparison_video.py --streams /tmp/streams --satellite 16 --out compare.mp4

The panes must stay in step, which is the whole reason this does not simply hand
ffmpeg three directories: a variant that produced no composite for a timestamp -- what
the current pipeline does past its gap limit -- would otherwise shorten that pane and
slide it out of alignment with the others.  Every stream is emitted against the same
canonical timeline, with black where a composite is missing, so frame N is the same
instant in all three panes.  Those black runs are not padding; they are what the
current pipeline actually shows.
"""

import argparse
import os
import re
import subprocess
import sys

import numpy as np
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from ffmpeg_video import find_ffmpeg
from suvi import paths, vfs

VARIANTS = ("pristine", "today", "new")
LABELS = {
    "pristine": "PRISTINE (ground truth)",
    "today": "TODAY (hold-last + black frames)",
    "new": "NEW (disc+header, optical flow)",
}
_STAMP = re.compile(r"Composite-(\d+)\.jpg$")


def stream_frames(directory):
    """{timestamp: path} for one rendered stream."""
    frames = {}
    try:
        entries = os.listdir(directory)
    except OSError:
        return frames
    for entry in entries:
        match = _STAMP.search(entry)
        if match:
            frames[int(match.group(1))] = os.path.join(directory, entry)
    return frames


def encode_stream(ffmpeg, frames, timeline, out_path, framerate, crf, reliever=None):
    """Encode one pane, emitting a black frame wherever the variant has nothing.

    Reads every composite in the pane, so it hands file handles back as it goes --
    three panes of a week-long window is over 7,000 files, enough on its own to
    exhaust the mount and take unrelated software down with it.
    """
    command = (
        f'{ffmpeg} -y -f image2pipe -framerate {framerate} -i - '
        f'-c:v libx264 -crf {crf} -preset veryfast -pix_fmt yuv420p "{out_path}"'
    )
    process = subprocess.Popen(command, shell=True, stdin=subprocess.PIPE,
                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    blank = None
    present = 0
    for when in timeline:
        path = frames.get(when)
        if path:
            with open(path, "rb") as handle:
                process.stdin.write(handle.read())
            present += 1
            if reliever is not None:
                reliever.tick()
        else:
            if blank is None:
                size = Image.open(next(iter(frames.values()))).size if frames else (1920, 1080)
                blank = Image.fromarray(np.zeros((size[1], size[0], 3), dtype="uint8"))
            blank.save(process.stdin, "jpeg", quality=95)
    process.stdin.close()
    process.wait()
    return present


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--streams", required=True, help="root written by bench.py render")
    parser.add_argument("--satellite", type=int, default=16)
    parser.add_argument("--out", required=True)
    parser.add_argument("--framerate", type=int, default=60)
    parser.add_argument("--crf", type=int, default=18)
    parser.add_argument("--keep-panes", action="store_true",
                        help="keep the intermediate per-pane videos")
    args = parser.parse_args(argv)

    ffmpeg = find_ffmpeg()
    streams = {}
    for variant in VARIANTS:
        directory = os.path.join(args.streams, variant, f"goes{args.satellite}")
        streams[variant] = stream_frames(directory)
        print(f"  {variant:>8}: {len(streams[variant])} composites in {directory}")
    if not any(streams.values()):
        print("No composites found; run bench.py render first.")
        return 1

    # One canonical timeline across every variant, so the panes stay in step.
    everything = set()
    for frames in streams.values():
        everything.update(frames)
    timeline = list(range(min(everything), max(everything) + paths.CADENCE, paths.CADENCE))
    print(f"  timeline: {len(timeline)} slots "
          f"({len(timeline) / args.framerate:.0f}s per pane at {args.framerate}fps)")

    panes = []
    reliever = vfs.Reliever(label="encode")
    for variant in VARIANTS:
        pane = f"{os.path.splitext(args.out)[0]}.{variant}.mp4"
        present = encode_stream(
            ffmpeg, streams[variant], timeline, pane, args.framerate, args.crf, reliever
        )
        missing = len(timeline) - present
        print(f"  encoded {variant}: {present} frames"
              + (f", {missing} black ({missing / len(timeline):.0%})" if missing else ""))
        panes.append(pane)

    inputs = " ".join(f'-i "{p}"' for p in panes)
    font = os.path.join(os.path.dirname(os.path.abspath(__file__)), "OpenSans-Regular.ttf")
    labelled = []
    for index, variant in enumerate(VARIANTS):
        labelled.append(
            f"[{index}:v]scale=960:-2,"
            f"drawtext=fontfile='{font}':text='{LABELS[variant]}':"
            f"x=12:y=12:fontsize=22:fontcolor=yellow:box=1:boxcolor=black@0.5[v{index}]"
        )
    graph = ";".join(labelled) + ";[v0][v1][v2]hstack=inputs=3[out]"
    command = (
        f'{ffmpeg} -y {inputs} -filter_complex "{graph}" -map "[out]" '
        f'-c:v libx264 -crf {args.crf} -preset veryfast -pix_fmt yuv420p "{args.out}"'
    )
    print("  stacking panes...")
    result = subprocess.run(command, shell=True, stdout=subprocess.DEVNULL,
                            stderr=subprocess.PIPE)
    if result.returncode != 0:
        print(result.stderr.decode()[-1500:])
        return 1

    if not args.keep_panes:
        for pane in panes:
            try:
                os.remove(pane)
            except OSError:
                pass
    reliever.finish()
    size = os.path.getsize(args.out) / 1024 / 1024
    print(f"Wrote {args.out} ({size:.0f} MB)")
    return 0


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