305 lines
13 KiB
Python
305 lines
13 KiB
Python
#!/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
|
|
|
|
`render_real.py` streams on real data work the same way (its ``--filler``
|
|
directories are variants), with ``@root`` per pane, ``--labels`` to name each
|
|
detection and filler, and ``--periods`` to show sample weeks back to back.
|
|
|
|
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 datetime
|
|
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
|
|
|
|
LABELS = {
|
|
"pristine": "PRISTINE (ground truth)",
|
|
"today": "TODAY (hold-last + black frames)",
|
|
"new": "OPTICAL FLOW (disc+header detection)",
|
|
"optical_flow": "OPTICAL FLOW",
|
|
"learned": "LEARNED (two-satellite aligned fusion)",
|
|
"damaged": "INPUT (gaps and damage, unrepaired)",
|
|
}
|
|
_STAMP = re.compile(r"Composite-(\d+)\.jpg$")
|
|
#: What a ``--labels`` entry may contain. The label lands inside a drawtext
|
|
#: option and a shell command line, where quotes, colons, backslashes and the
|
|
#: like would end the option or the command.
|
|
_LABEL = re.compile(r"[A-Za-z0-9 ._+()|/-]+")
|
|
|
|
|
|
def parse_labels(spec, count):
|
|
"""'a,b,...' -> one on-screen label per pane, or None for the defaults."""
|
|
if spec is None:
|
|
return None
|
|
labels = tuple(label.strip() for label in spec.split(","))
|
|
if len(labels) != count:
|
|
raise SystemExit(f"--labels names {len(labels)} panes; there are {count}")
|
|
for label in labels:
|
|
if not _LABEL.fullmatch(label):
|
|
raise SystemExit(f"label {label!r}: use letters, digits, spaces and ._+()|/-")
|
|
return labels
|
|
|
|
|
|
def parse_panes(spec, default_satellite, default_root=None):
|
|
"""'variant[:satellite][@root],...' -> ((variant, satellite, root), ...).
|
|
|
|
``@root`` reads that pane from a different ``bench.py render`` output
|
|
directory. A new model only needs its own ``learned`` stream rendered; the
|
|
damaged inputs and the optical-flow repair of the same case already exist
|
|
elsewhere, and this share supports neither symlinks nor cheap copies of a
|
|
few thousand multi-megabyte composites.
|
|
"""
|
|
panes = []
|
|
for entry in spec.split(","):
|
|
entry = entry.strip()
|
|
if not entry:
|
|
continue
|
|
entry, _, root = entry.partition("@")
|
|
variant, _, satellite = entry.partition(":")
|
|
if variant not in LABELS:
|
|
raise SystemExit(f"unknown variant {variant!r}; known: {sorted(LABELS)}")
|
|
panes.append((variant, int(satellite) if satellite else default_satellite,
|
|
root or default_root))
|
|
return tuple(panes)
|
|
|
|
|
|
def pane_label(variant, satellite, root, default_root):
|
|
"""The on-screen label; a pane read from its own root names that root, so two
|
|
models' outputs can never be mistaken for one another."""
|
|
text = f"{LABELS[variant]} | GOES-{satellite}"
|
|
if root and root != default_root:
|
|
text += f" | {os.path.basename(os.path.normpath(root))}"
|
|
return text
|
|
|
|
|
|
def filter_graph(panes, default_root, font, labels=None):
|
|
"""The ffmpeg filter graph that labels each pane and lays them out as [out].
|
|
|
|
One pane is a single product stream rather than a comparison, so it keeps its
|
|
full resolution and is not stacked: ffmpeg's stack filters need two inputs.
|
|
`labels`, from :func:`parse_labels`, replaces the default pane labels.
|
|
|
|
Labels sit bottom centre, in the empty corona below the disc. At the top
|
|
they covered the composite's own title and timestamp, which is how a frame
|
|
seen in the video is found again.
|
|
"""
|
|
scale = "scale=960:-2," if len(panes) > 1 else ""
|
|
labelled = []
|
|
for index, (variant, satellite, root) in enumerate(panes):
|
|
text = (labels[index] if labels
|
|
else pane_label(variant, satellite, root, default_root))
|
|
labelled.append(
|
|
f"[{index}:v]{scale}"
|
|
f"drawtext=fontfile='{font}':text='{text}':"
|
|
f"x=(w-text_w)/2:y=h-text_h-12:fontsize=22:fontcolor=yellow:"
|
|
f"box=1:boxcolor=black@0.5[v{index}]"
|
|
)
|
|
taps = "".join(f"[v{i}]" for i in range(len(panes)))
|
|
if len(panes) == 1:
|
|
stack = f"{taps}null[out]"
|
|
elif len(panes) == 4:
|
|
# Row-major 2x2: panes[0] top-left, [1] top-right, [2] bottom-left,
|
|
# [3] bottom-right.
|
|
stack = f"{taps}xstack=inputs=4:layout=0_0|w0_0|0_h0|w0_h0[out]"
|
|
else:
|
|
stack = f"{taps}hstack=inputs={len(panes)}[out]"
|
|
return ";".join(labelled) + ";" + stack
|
|
|
|
|
|
def parse_periods(spec):
|
|
"""'2024-01-01..2024-01-08,...' -> ((t0, t1), ...) in UTC seconds, ends
|
|
exclusive, in time order. Overlaps are refused: they would show the same
|
|
slots twice."""
|
|
periods = []
|
|
for entry in spec.split(","):
|
|
start, separator, end = entry.strip().partition("..")
|
|
try:
|
|
t0, t1 = (int(datetime.datetime.strptime(text, "%Y-%m-%d").replace(
|
|
tzinfo=datetime.timezone.utc).timestamp()) for text in (start, end))
|
|
except ValueError:
|
|
separator = ""
|
|
if not separator:
|
|
raise SystemExit(f"period {entry!r}: expected YYYY-MM-DD..YYYY-MM-DD")
|
|
if t1 <= t0:
|
|
raise SystemExit(f"period {entry!r} ends before it starts")
|
|
periods.append((t0, t1))
|
|
periods.sort()
|
|
for (_, end), (start, _) in zip(periods, periods[1:]):
|
|
if start < end:
|
|
raise SystemExit(f"periods overlap: {spec}")
|
|
return tuple(periods)
|
|
|
|
|
|
def build_timeline(streams, periods=None):
|
|
"""The slots every pane is emitted against.
|
|
|
|
By default every cadence slot from the first composite in any pane to the
|
|
last. With `periods`, only the slots inside them, back to back: a pane read
|
|
from a whole-year render would otherwise stretch three sample weeks into a
|
|
year of mostly black video.
|
|
"""
|
|
if periods:
|
|
return [when for t0, t1 in periods for when in range(t0, t1, paths.CADENCE)]
|
|
everything = set()
|
|
for frames in streams.values():
|
|
everything.update(frames)
|
|
if not everything:
|
|
return []
|
|
return list(range(min(everything), max(everything) + paths.CADENCE, paths.CADENCE))
|
|
|
|
|
|
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):
|
|
"""Encode one pane, emitting a black frame wherever the variant has nothing.
|
|
|
|
Reads every composite in the pane -- three panes of a week-long window is over
|
|
7,000 files -- which is enough to fill the guest's inode cache on its own. Keeping
|
|
that below the mount's ceiling is the reclaim watchdog's job, not this loop's; see
|
|
reclaim.py.
|
|
"""
|
|
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
|
|
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("--variants", default="damaged:18,new,damaged,learned",
|
|
help="comma-separated pane order, each "
|
|
"'variant[:satellite][@root]' (default satellite "
|
|
"from --satellite, default root --streams). Four "
|
|
"panes make a 2x2 grid (row-major); any other count "
|
|
"stacks horizontally")
|
|
parser.add_argument("--labels", default=None,
|
|
help="comma-separated on-screen label per pane, replacing "
|
|
"the defaults (letters, digits, spaces and ._+()|/-)")
|
|
parser.add_argument("--periods", default=None,
|
|
help="comma-separated UTC date ranges START..END (end "
|
|
"exclusive) to show, back to back (default: everything "
|
|
"from the first composite to the last)")
|
|
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()
|
|
panes_wanted = parse_panes(args.variants, args.satellite, args.streams)
|
|
labels = parse_labels(args.labels, len(panes_wanted))
|
|
periods = parse_periods(args.periods) if args.periods else None
|
|
streams = {}
|
|
for key in panes_wanted:
|
|
variant, satellite, root = key
|
|
directory = os.path.join(root, variant, f"goes{satellite}")
|
|
streams[key] = stream_frames(directory)
|
|
print(f" {variant:>8} g{satellite}: "
|
|
f"{len(streams[key])} composites in {directory}")
|
|
# One canonical timeline across every variant, so the panes stay in step.
|
|
timeline = build_timeline(streams, periods)
|
|
if not any(when in frames for frames in streams.values() for when in timeline):
|
|
print("No composites found; run bench.py render first.")
|
|
return 1
|
|
print(f" timeline: {len(timeline)} slots "
|
|
f"({len(timeline) / args.framerate:.0f}s per pane at {args.framerate}fps)")
|
|
|
|
panes = []
|
|
for index, key in enumerate(panes_wanted):
|
|
variant, satellite, _ = key
|
|
pane = f"{os.path.splitext(args.out)[0]}.{index}-{variant}-g{satellite}.mp4"
|
|
present = encode_stream(
|
|
ffmpeg, streams[key], timeline, pane, args.framerate, args.crf
|
|
)
|
|
missing = len(timeline) - present
|
|
print(f" encoded {variant} g{satellite}: {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")
|
|
graph = filter_graph(panes_wanted, args.streams, font, labels)
|
|
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
|
|
size = os.path.getsize(args.out) / 1024 / 1024
|
|
print(f"Wrote {args.out} ({size:.0f} MB)")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|