"""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) 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"))