added 5 sensors multi-stream support, 5 finger reach > 100 fps stable using .mp4 video

2026-09-13 12:18:27 +08:00
parent 0d5d8ff70d
commit 514874e2f4
10 changed files with 1900 additions and 7 deletions
@@ -0,0 +1,281 @@
+"""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"))