ykcheungab
282 lines
8.8 KiB
Python
282 lines
8.8 KiB
Python
"""StreamRecorder / mosaic_videos: off-hot-path arrow overlay export."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import time
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from orisys import FrameSnapshot, StreamRecorder, mosaic_videos
|
|
|
|
|
|
_DEFAULT = object()
|
|
|
|
|
|
def make_snapshot(seq: int = 0, *, size: int = 64, image=_DEFAULT, flow=_DEFAULT) -> FrameSnapshot:
|
|
if image is _DEFAULT:
|
|
image = np.full((size, size), 128, dtype=np.uint8)
|
|
image[8:-8, 8:-8] = 200
|
|
if flow is _DEFAULT:
|
|
flow = np.zeros((size, size, 2), dtype=np.float32)
|
|
flow[..., 0] = 2.0
|
|
flow[..., 1] = 1.0
|
|
return FrameSnapshot(
|
|
seq=seq,
|
|
timestamp=time.time(),
|
|
age_ms=5.0,
|
|
fps=100.0,
|
|
unit="raw",
|
|
normal_force=1.5,
|
|
shear_x=0.25,
|
|
shear_y=-0.5,
|
|
normal_force_raw=1.5,
|
|
shear_x_raw=0.25,
|
|
shear_y_raw=-0.5,
|
|
torque=0.01,
|
|
image=image,
|
|
flow=flow,
|
|
)
|
|
|
|
|
|
def read_video(path):
|
|
cap = cv2.VideoCapture(path)
|
|
try:
|
|
assert cap.isOpened(), f"cannot reopen {path}"
|
|
frames = []
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
frames.append(frame)
|
|
return frames, (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
|
|
int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
|
|
finally:
|
|
cap.release()
|
|
|
|
|
|
def _diff_count(frame, expected, tol=20):
|
|
"""Pixels differing from ``expected`` by more than the codec's noise floor."""
|
|
d = np.abs(frame.astype(np.int16) - expected.astype(np.int16)).max(axis=2)
|
|
return int((d > tol).sum())
|
|
|
|
|
|
def _near(pixel, value, tol=14):
|
|
"""Codec round-trip shifts solid greys by a few levels (YUV range)."""
|
|
return all(abs(int(c) - value) <= tol for c in np.ravel(pixel))
|
|
|
|
|
|
# --------------------------------------------------------------------- recorder
|
|
def test_recorder_writes_every_frame(tmp_path):
|
|
path = str(tmp_path / "s0.mp4")
|
|
# Depth >= frame count so the drop-oldest policy provably does not engage;
|
|
# exactness of the writer itself is what this test is about.
|
|
rec = StreamRecorder(path, fps=25.0, grid_spacing=8, overlay_values=False,
|
|
queue_depth=64)
|
|
for i in range(12):
|
|
rec.on_frame(make_snapshot(i))
|
|
stats = rec.close()
|
|
|
|
assert stats["seen"] == 12
|
|
assert stats["written"] == 12
|
|
assert stats["dropped"] == 0
|
|
assert stats["skipped"] == 0
|
|
assert "error" not in stats
|
|
assert stats["thread_alive"] is False
|
|
|
|
frames, size = read_video(path)
|
|
assert len(frames) == 12
|
|
assert size == (64, 64)
|
|
|
|
|
|
def test_recorder_every_skips_frames(tmp_path):
|
|
path = str(tmp_path / "s0.mp4")
|
|
rec = StreamRecorder(path, every=3, grid_spacing=8, overlay_values=False,
|
|
queue_depth=64)
|
|
for i in range(12):
|
|
rec.on_frame(make_snapshot(i))
|
|
stats = rec.close()
|
|
|
|
assert stats["seen"] == 12
|
|
assert stats["written"] == 4
|
|
assert stats["skipped"] == 8
|
|
frames, _ = read_video(path)
|
|
assert len(frames) == 4
|
|
|
|
|
|
def test_on_frame_never_blocks_and_drops_oldest(tmp_path, monkeypatch):
|
|
"""A slow writer must cost frames, never pipeline throughput."""
|
|
path = str(tmp_path / "s0.mp4")
|
|
rec = StreamRecorder(path, queue_depth=2, grid_spacing=8, overlay_values=False)
|
|
|
|
original = rec._render
|
|
|
|
def slow_render(snap):
|
|
time.sleep(0.05)
|
|
return original(snap)
|
|
|
|
rec._render = slow_render
|
|
|
|
t0 = time.perf_counter()
|
|
for i in range(20):
|
|
rec.on_frame(make_snapshot(i))
|
|
ingest = time.perf_counter() - t0
|
|
stats = rec.close(timeout=30.0)
|
|
|
|
# Ingesting 20 frames must not wait on ~1s of encoder work.
|
|
assert ingest < 0.1, f"on_frame blocked for {ingest * 1000:.1f} ms"
|
|
assert stats["dropped"] > 0
|
|
assert stats["written"] + stats["dropped"] + stats["skipped"] == 20
|
|
assert stats["written"] < 20
|
|
|
|
|
|
def test_recorder_reports_missing_arrays(tmp_path):
|
|
path = str(tmp_path / "s0.mp4")
|
|
rec = StreamRecorder(path, grid_spacing=8, overlay_values=False)
|
|
rec.on_frame(make_snapshot(0, image=None))
|
|
stats = rec.close()
|
|
|
|
assert stats["written"] == 0
|
|
assert "error" in stats
|
|
assert 'include=("image", "flow")' in stats["error"]
|
|
|
|
|
|
def test_recorder_context_manager(tmp_path):
|
|
path = str(tmp_path / "s0.mp4")
|
|
with StreamRecorder(path, grid_spacing=8, overlay_values=False) as rec:
|
|
rec.on_frame(make_snapshot(0))
|
|
assert os.path.isfile(path)
|
|
assert rec._writer is None
|
|
|
|
|
|
def test_valid_mask_callable_resolved_per_frame(tmp_path):
|
|
calls = []
|
|
|
|
def mask():
|
|
calls.append(1)
|
|
return None
|
|
|
|
path = str(tmp_path / "s0.mp4")
|
|
rec = StreamRecorder(path, grid_spacing=8, overlay_values=False, valid_mask=mask)
|
|
for i in range(5):
|
|
rec.on_frame(make_snapshot(i))
|
|
stats = rec.close()
|
|
|
|
assert stats["written"] == 5
|
|
assert len(calls) == 5
|
|
|
|
|
|
def test_all_false_mask_suppresses_arrows(tmp_path):
|
|
"""With every grid point masked out the overlay is just the grey frame."""
|
|
path = str(tmp_path / "s0.mp4")
|
|
snap = make_snapshot(0)
|
|
rec = StreamRecorder(path, grid_spacing=8, overlay_values=False,
|
|
valid_mask=np.zeros((64, 64), dtype=bool))
|
|
rec.on_frame(snap)
|
|
rec.close()
|
|
|
|
frames, _ = read_video(path)
|
|
expected = cv2.cvtColor(snap.image, cv2.COLOR_GRAY2BGR)
|
|
assert _diff_count(frames[0], expected) == 0
|
|
|
|
|
|
def test_arrows_change_pixels_without_mask(tmp_path):
|
|
path = str(tmp_path / "s0.mp4")
|
|
snap = make_snapshot(0)
|
|
rec = StreamRecorder(path, grid_spacing=8, overlay_values=False)
|
|
rec.on_frame(snap)
|
|
rec.close()
|
|
|
|
frames, _ = read_video(path)
|
|
expected = cv2.cvtColor(snap.image, cv2.COLOR_GRAY2BGR)
|
|
assert _diff_count(frames[0], expected) > 100
|
|
|
|
|
|
def test_close_without_frames_is_safe(tmp_path):
|
|
rec = StreamRecorder(str(tmp_path / "empty.mp4"), grid_spacing=8)
|
|
stats = rec.close()
|
|
assert stats["written"] == 0
|
|
assert "error" not in stats
|
|
|
|
|
|
def test_stats_report_requested_geometry(tmp_path):
|
|
rec = StreamRecorder(str(tmp_path / "s0.mp4"), fps=12.0, fourcc="mp4v")
|
|
assert rec.stats()["fps"] == 12.0
|
|
assert rec.stats()["fourcc"] == "mp4v"
|
|
assert rec.stats()["path"].endswith("s0.mp4")
|
|
|
|
|
|
# ---------------------------------------------------------------------- mosaic
|
|
def _write_solid(path, value, frames=4, size=(16, 12), fps=10.0):
|
|
w, h = size
|
|
writer = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h), True)
|
|
assert writer.isOpened()
|
|
try:
|
|
for _ in range(frames):
|
|
writer.write(np.full((h, w, 3), value, dtype=np.uint8))
|
|
finally:
|
|
writer.release()
|
|
return path
|
|
|
|
|
|
def test_mosaic_grid_geometry(tmp_path):
|
|
srcs = [_write_solid(str(tmp_path / f"s{i}.mp4"), 40 + 40 * i) for i in range(3)]
|
|
out = str(tmp_path / "mosaic.mp4")
|
|
info = mosaic_videos(srcs, out, cols=2, gap=4, fps=10.0)
|
|
|
|
assert info["grid"] == (2, 2)
|
|
# 2 columns x 16 + 1 gap, 2 rows x 12 + 1 gap
|
|
assert info["size"] == (2 * 16 + 4, 2 * 12 + 4)
|
|
assert info["frames"] == 4
|
|
frames, size = read_video(out)
|
|
assert len(frames) == 4
|
|
assert size == info["size"]
|
|
# the empty bottom-right cell stays black
|
|
assert _near(frames[0][-1, -1], 0, tol=12)
|
|
# top-left tile carries stream 0's colour
|
|
assert _near(frames[0][0, 0], 40)
|
|
|
|
|
|
def test_mosaic_holds_last_frame_on_short_input(tmp_path):
|
|
a = _write_solid(str(tmp_path / "a.mp4"), 200, frames=5)
|
|
b = _write_solid(str(tmp_path / "b.mp4"), 50, frames=3)
|
|
out = str(tmp_path / "mosaic.mp4")
|
|
info = mosaic_videos([a, b], out, cols=2, gap=0, fps=10.0)
|
|
|
|
assert info["frames"] == 5
|
|
frames, size = read_video(out)
|
|
assert len(frames) == 5
|
|
# b's 4th frame is a repeat of its 3rd, so the right half is still solid 50
|
|
assert _near(frames[-1][0, size[0] // 2], 50)
|
|
|
|
|
|
def test_mosaic_labels_drawn(tmp_path):
|
|
srcs = [_write_solid(str(tmp_path / f"s{i}.mp4"), 30) for i in range(2)]
|
|
out = str(tmp_path / "mosaic.mp4")
|
|
plain = str(tmp_path / "plain.mp4")
|
|
mosaic_videos(srcs, plain, cols=2, gap=0, fps=10.0)
|
|
mosaic_videos(srcs, out, cols=2, gap=0, fps=10.0, labels=["#0", "#1"])
|
|
|
|
with_label, _ = read_video(out)
|
|
without, _ = read_video(plain)
|
|
assert _diff_count(with_label[0], without[0]) > 0
|
|
|
|
|
|
def test_mosaic_scale_downscales(tmp_path):
|
|
srcs = [_write_solid(str(tmp_path / f"s{i}.mp4"), 30, size=(16, 12)) for i in range(4)]
|
|
out = str(tmp_path / "mosaic.mp4")
|
|
info = mosaic_videos(srcs, out, cols=2, gap=2, scale=0.5, fps=10.0)
|
|
assert info["size"] == (8 * 2 + 2, 6 * 2 + 2)
|
|
|
|
|
|
def test_mosaic_rejects_missing_input(tmp_path):
|
|
with pytest.raises(FileNotFoundError):
|
|
mosaic_videos([str(tmp_path / "nope.mp4")], str(tmp_path / "out.mp4"))
|
|
|
|
|
|
def test_mosaic_rejects_empty_input(tmp_path):
|
|
with pytest.raises(ValueError):
|
|
mosaic_videos([], str(tmp_path / "out.mp4"))
|