diff --git a/.gitignore b/.gitignore index 4aaeb7d..d440858 100644 --- a/.gitignore +++ b/.gitignore @@ -42,4 +42,7 @@ __pycache__/ # Legacy / scratch folders /old_orisys/ -/tools/calibration/ \ No newline at end of file +/tools/calibration/ + +# Scratch output of examples/ (multi_stream.py's default --out-dir) +/outputs/ \ No newline at end of file diff --git a/examples/multi_stream.py b/examples/multi_stream.py new file mode 100644 index 0000000..7944929 --- /dev/null +++ b/examples/multi_stream.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python3 +"""Multi-sensor real-time example: N streams, one process each. + +This is the deployment shape the multi-sensor throughput numbers assume. Each +stream owns a process and pipelines the per-frame work into three concurrent +stages (capture+unroll -> DIS flow -> HHD+readout), so a stream's rate is bounded +by its slowest stage instead of the sum of all three. + + # five finger streams from a video file (benchmark / smoke test) + python examples/multi_stream.py --streams 5 -c finger -v data/dd02-0703/1.mp4 + + # five live cameras, camera ids 0..4 + python examples/multi_stream.py --streams 5 -c finger --camera-base 0 + + # single box sensor + python examples/multi_stream.py --streams 1 -c box -v data/long.mp4 + + # five streams, additionally recording an arrow-overlay video per stream + # plus one 5-up mosaic, without stalling the pipeline + python examples/multi_stream.py --streams 5 -c finger -v data/dd02-0703/1.mp4 \ + --seconds 10 --record outputs/arrows --mosaic + +Notes +----- +* One process per stream. Threads inside a single process share the GIL and, + measured on a 9950X, five finger streams reach ~150 FPS aggregate as threads + versus ~626 FPS as processes. +* ``cv2.setNumThreads(2)`` per process: at five streams that beats 1, 4 and 8. +* Run one stream per process even when the sensor count is small; the process + boundary also isolates a crashed camera from the rest. +* ``--record`` draws the DIS flow with ``orisys.util.draw_arrows`` (plus the + contact overlay with ``--contact``) and encodes it with ``cv2.VideoWriter``. + Both happen on a dedicated thread fed by a bounded drop-oldest queue, so the + pipeline never waits on the encoder; frames the writer cannot keep up with are + counted in the ``rec`` column as ``written/dropped``. Keep ``--record`` for + verification runs - the encoder still competes for CPU with the pipeline. +""" + +from __future__ import annotations + +import argparse +import json +import multiprocessing as mp +import os +import statistics as st +import sys +import time + +os.environ.setdefault("OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS", "0") + +import orisys +from orisys import StreamPipeline, StreamRecorder, configs + + +def _make_stream(args, index): + """Return (sensor, produce, close) for one stream.""" + source = args.video + if args.camera_base is not None: + source = int(args.camera_base) + index + if source is None: + raise ValueError("pass --video PATH or --camera-base ID") + + cfg = configs.load_config(args.config) + if str(cfg.get("type", "")).lower() == "single_camera": + from orisys.finger import FingerRuntime + + runtime = FingerRuntime.open(config=args.config, video=source, + verbose=False) + return runtime.tactile.raw_sensor, _finger_produce(runtime), runtime.close + + sensor = orisys.Sensor(source, config_name=args.config, isstitch=True, + cuda=False, backend="cpu", verbose=False) + + def produce(): + return sensor.get_img() + + return sensor, produce, sensor.disconnect + + +def _finger_produce(runtime): + def produce(): + ok, frame = runtime.grab_and_process(color=False) + return frame if ok else None + + return produce + + +def worker(args, index, out_path): + """One stream: own process, three-stage pipeline.""" + import cv2 + + cv2.setNumThreads(int(args.threads)) + sensor, produce, close = _make_stream(args, index) + + frames = [] + stamps = [] + started = time.perf_counter() + + def collect(snap): + frames.append(snap.age_ms) + stamps.append(time.perf_counter()) + + # Recording is entirely off the hot path: ``on_frame`` only pushes onto a + # bounded drop-oldest queue, and a separate thread does the drawing and the + # H.264/MPEG-4 encoding. The pipeline never blocks on the encoder. + recorder = None + on_frame = collect + if args.record: + recorder = StreamRecorder( + os.path.join(args.record, f"stream_{index}.mp4"), + fps=args.record_fps, + every=args.record_every, + fourcc=args.record_fourcc, + valid_mask=lambda: getattr(sensor, "valid_mask", None), + grid_spacing=args.record_grid, + threshold=args.record_threshold, + arrow_scale=args.record_arrow_scale, + draw_contact=args.contact, + label=f"#{index}", + verbose=True, + ) + + def on_frame(snap, _collect=collect, _rec=recorder): + _collect(snap) + _rec.on_frame(snap) + + if args.no_arrays: + include = () + elif recorder is not None: + include = ("image", "flow") # what the arrow overlay needs + else: + include = ("image",) + + pipe = StreamPipeline(sensor, produce, on_frame=on_frame, + include=include, contact=args.contact, depth=args.depth) + pipe.start() + record_stats = None + try: + deadline = started + args.seconds if args.seconds else None + while True: + if not pipe.running: + break + if args.frames and pipe.stats.published >= args.frames: + break + if deadline is not None and time.perf_counter() >= deadline: + break + time.sleep(0.002) + finally: + try: + pipe.stop() + except BaseException as exc: # surface a stage failure + print(f"[stream {index}] pipeline error: {type(exc).__name__}: {exc}", + file=sys.stderr) + raise + finally: + if recorder is not None: + record_stats = recorder.close() + close() + + wall = time.perf_counter() - started + # Steady-state rate: skip the first frames, which pay the JIT / HHD-cache / + # DIS-reference-frame warm-up and would drag a short run's average down. + warm = min(30, max(0, len(stamps) - 2)) + span = (stamps[-1] - stamps[warm]) if len(stamps) > warm + 1 else 0.0 + steady = (len(stamps) - 1 - warm) / span if span > 0 else 0.0 + payload = { + "index": index, + "published": pipe.stats.published, + "wall": wall, + "fps": pipe.stats.published / wall if wall else 0.0, + "steady_fps": steady, + "stats": pipe.stats.as_dict(), + "age_p95_ms": float(st.quantiles(frames, n=20)[18]) if len(frames) >= 20 else 0.0, + "record": record_stats, + } + with open(out_path, "w") as fh: + json.dump(payload, fh) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--streams", "-n", type=int, default=5) + ap.add_argument("--config", "-c", default="finger") + ap.add_argument("--video", "-v", default=None, help="video file shared by all streams") + ap.add_argument("--camera-base", type=int, default=None, + help="first camera id; stream i uses base + i") + ap.add_argument("--frames", type=int, default=0, help="stop after N frames (0 = until EOF)") + ap.add_argument("--seconds", type=float, default=10.0, help="max wall time per stream") + ap.add_argument("--threads", type=int, default=2, help="cv2.setNumThreads per stream") + ap.add_argument("--depth", type=int, default=2, help="queue depth between stages") + ap.add_argument("--contact", action="store_true", help="also run compute_contact") + ap.add_argument("--no-arrays", action="store_true", + help="scalars only (fastest; skips the snapshot image copy)") + ap.add_argument("--out-dir", default=None) + ap.add_argument("--record", default=None, metavar="DIR", + help="write one arrow-overlay mp4 per stream into DIR " + "(drawing+encoding run on a separate thread, off the " + "pipeline's critical path)") + ap.add_argument("--record-fps", type=float, default=30.0, + help="playback fps of the recorded files") + ap.add_argument("--record-every", type=int, default=1, + help="keep one frame in N (raise it if the writer falls behind)") + ap.add_argument("--record-fourcc", default="mp4v", + help="mp4v (built in) or MJPG/avi for maximum compatibility") + ap.add_argument("--record-grid", type=int, default=14, help="arrow grid spacing in px") + ap.add_argument("--record-threshold", type=float, default=0.1, + help="displacement (px) below which a dot is drawn instead of an arrow") + ap.add_argument("--record-arrow-scale", type=float, default=0.15, + help="draw_arrows divides the displacement by this; 0.15 magnifies " + "the finger's ~1-3 px field ~6.7x so directions are legible") + ap.add_argument("--mosaic", action="store_true", + help="after the run, tile the recordings into one grid video") + ap.add_argument("--mosaic-cols", type=int, default=0, help="mosaic columns (0 = square-ish)") + ap.add_argument("--mosaic-scale", type=float, default=1.0, help="mosaic downscale factor") + args = ap.parse_args() + + if args.record and args.no_arrays: + ap.error("--record needs the frame and flow arrays; drop --no-arrays") + if args.mosaic and not args.record: + ap.error("--mosaic needs --record DIR") + + if args.record: + args.record = os.path.abspath(args.record) + os.makedirs(args.record, exist_ok=True) + + out_dir = args.out_dir or os.path.join(os.path.dirname(os.path.abspath(__file__)), + "..", "outputs", "multi_stream") + out_dir = os.path.abspath(out_dir) + os.makedirs(out_dir, exist_ok=True) + + ctx = mp.get_context("spawn") + procs, paths = [], [] + for i in range(args.streams): + path = os.path.join(out_dir, f"stream_{i}.json") + paths.append(path) + procs.append(ctx.Process(target=worker, args=(args, i, path))) + for p in procs: + p.start() + for p in procs: + p.join() + + rows = [] + for path in paths: + try: + with open(path) as fh: + rows.append(json.load(fh)) + except Exception: + pass + if not rows: + print("no stream produced results", file=sys.stderr) + return 1 + + wall = max(r["wall"] for r in rows) + total = sum(r["published"] for r in rows) + steady_rows = [r for r in rows if r.get("steady_fps")] + steady = st.mean(r["steady_fps"] for r in steady_rows) if steady_rows else 0.0 + print(f"\nstreams={len(rows)} frames={total} wall={wall:.2f}s") + print(f"aggregate throughput : {total / wall:.1f} FPS (incl. warm-up)") + print(f"per-stream throughput : {total / wall / len(rows):.1f} FPS (incl. warm-up)") + print(f"per-stream steady-state : {steady:.1f} FPS") + recording = args.record is not None + print(f"\n{'stream':>7}{'published':>11}{'FPS':>9}{'steady':>9}{'produce':>10}" + f"{'flow':>9}{'hhd':>9}{'age':>9}{'age p95':>10}" + + (f"{'rec':>13}" if recording else "")) + for r in sorted(rows, key=lambda x: x["index"]): + s = r["stats"] + rec = "" + if recording: + rs = r.get("record") or {} + rec = f"{rs.get('written', 0)}/{rs.get('dropped', 0)}" + if rs.get("error"): + rec = "ERR" + print(f"{r['index']:>7}{r['published']:>11}{r['fps']:>9.1f}" + f"{r.get('steady_fps', 0.0):>9.1f}" + f"{s['produce_ms']:>10.2f}{s['flow_ms']:>9.2f}{s['hhd_ms']:>9.2f}" + f"{s['age_ms']:>9.2f}{r['age_p95_ms']:>10.2f}" + + (f"{rec:>13}" if recording else "")) + if recording: + print("rec column = frames written / frames dropped by the writer thread") + for r in sorted(rows, key=lambda x: x["index"]): + rs = r.get("record") or {} + if rs.get("error"): + print(f" stream {r['index']}: {rs['error']}") + + if args.mosaic and recording: + streams = sorted(r["index"] for r in rows) + srcs = [os.path.join(args.record, f"stream_{i}.mp4") for i in streams] + missing = [p for p in srcs if not os.path.isfile(p)] + if missing: + print(f"mosaic skipped, missing {missing}", file=sys.stderr) + else: + out = os.path.join(args.record, "mosaic.mp4") + # No tile labels: each recorded frame already carries the "#N" HUD. + info = orisys.mosaic_videos(srcs, out, cols=args.mosaic_cols or None, + scale=args.mosaic_scale, + fps=args.record_fps) + print(f"\nmosaic -> {info['path']} {info['size'][0]}x{info['size'][1]} " + f"{info['grid'][0]}x{info['grid'][1]} grid {info['frames']} frames " + f"@ {info['fps']:g} fps") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/orisys/__init__.py b/src/orisys/__init__.py index 5e4f86b..b8773bd 100644 --- a/src/orisys/__init__.py +++ b/src/orisys/__init__.py @@ -15,6 +15,8 @@ from . import calibration from . import finger from . import signal_processing from . import capture +from .pipeline import FrameSnapshot, PipelineStats, StreamPipeline +from .export import StreamRecorder, mosaic_videos from .frame_processors import FFCFrameProcessor CompoundEyeSensor = Sensor @@ -44,6 +46,15 @@ __all__ = [ # capture helpers 'capture', + # pipelined multi-stream processing + 'StreamPipeline', + 'FrameSnapshot', + 'PipelineStats', + + # offline recording of arrow overlays + 'StreamRecorder', + 'mosaic_videos', + # force calibration 'calibration', diff --git a/src/orisys/export.py b/src/orisys/export.py new file mode 100644 index 0000000..98f577f --- /dev/null +++ b/src/orisys/export.py @@ -0,0 +1,365 @@ +"""Record tactile arrow overlays to video without stalling the pipeline. + +Drawing the arrow grid costs a few milliseconds per frame (hundreds of +``cv2.arrowedLine`` calls) and encoding costs a few more, so doing it inline +would cut a stream's throughput. Instead: + +* :meth:`StreamRecorder.on_frame` is designed to be the pipeline's ``on_frame`` + callback: it only performs a bounded, drop-oldest queue push, so the slowest + pipeline stage is unaffected and the run never blocks on the encoder; +* a dedicated writer thread does the drawing and the ``cv2.VideoWriter`` calls; +* when the writer falls behind, the newest frames win and the rest are counted + in :meth:`StreamRecorder.stats` so a recording gap is never silent. + +Typical use:: + + recorder = StreamRecorder("run/stream_0.mp4", label="#0", + valid_mask=lambda: runtime.valid_mask) + pipe = StreamPipeline(sensor, produce, include=("image", "flow"), + on_frame=recorder.on_frame) + pipe.start() + ... + print(recorder.close()) # written / dropped / fps + +Per-stream files are the natural unit because each stream is its own process; +:func:`mosaic_videos` tiles them into one grid afterwards, which costs nothing +during the run. + +Arrow parameters +---------------- +``draw_arrows`` divides the sampled displacement by ``arrow_scale``, so the +defaults assume the finger's UV domain, where the active-region displacement is +only ~1-3 px: ``grid_spacing=14, arrow_scale=0.15, threshold=0.1`` magnifies it +~6.7x, which is what makes the direction field legible instead of a scatter of +dots (the ``draw_arrows`` defaults of ``arrow_scale=1.0, threshold=5`` render the +raw pixel displacement and are tuned for large-motion scenes). Length saturates +at ``grid_spacing * margin`` px, so colour - not length - carries the magnitude. +A box/mini stream at 365x365 needs its own calibration of ``arrow_scale``. +""" + +from __future__ import annotations + +import os +import queue +import threading +import time +from typing import Any, Callable, Iterable, Sequence + +import cv2 +import numpy as np + +from .util import draw_arrows, draw_contact + +__all__ = ["StreamRecorder", "mosaic_videos"] + + +class StreamRecorder: + """Draw arrow overlays from snapshots and encode them to a video file.""" + + def __init__( + self, + path: str, + *, + fps: float = 30.0, + every: int = 1, + fourcc: str = "mp4v", + valid_mask: np.ndarray | Callable[[], np.ndarray | None] | None = None, + grid_spacing: int = 14, + threshold: float = 0.1, + arrow_scale: float = 0.15, + margin: float = 0.8, + draw_contact: bool = False, + label: str | None = None, + overlay_values: bool = True, + queue_depth: int = 8, + verbose: bool = False, + ) -> None: + self.path = str(path) + self.fps = max(1.0, float(fps)) + self.every = max(1, int(every)) + self.fourcc = fourcc + # A callable is resolved once per frame: the SDK rebuilds + # ``Sensor.valid_mask`` during ``get_img``, so holding an array taken + # before the first frame would capture ``None``. + self.valid_mask = valid_mask + self.grid_spacing = int(grid_spacing) + self.threshold = float(threshold) + self.arrow_scale = float(arrow_scale) + self.margin = float(margin) + self.draw_contact_flag = bool(draw_contact) + self.label = label + self.overlay_values = bool(overlay_values) + self.verbose = bool(verbose) + + self._queue: queue.Queue = queue.Queue(maxsize=max(1, int(queue_depth))) + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self._seen = 0 + self._written = 0 + self._dropped = 0 + self._skipped = 0 + self._error: BaseException | None = None + self._writer: cv2.VideoWriter | None = None + self._size: tuple[int, int] | None = None + + # ------------------------------------------------------------------ ingest + def on_frame(self, snap: Any) -> None: + """Pipeline callback: enqueue one snapshot. Never blocks.""" + self._seen += 1 + if self.every > 1 and (self._seen - 1) % self.every: + self._skipped += 1 + return + while True: + try: + self._queue.put_nowait(snap) + break + except queue.Full: + # Drop the oldest so the newest frame is always kept. + try: + self._queue.get_nowait() + self._dropped += 1 + except queue.Empty: + pass + self._ensure_thread() + + def _ensure_thread(self) -> None: + if self._thread is None: + self._thread = threading.Thread(target=self._loop, daemon=True, + name="orisys-recorder") + self._thread.start() + + # ------------------------------------------------------------------ writer + def _open(self, frame: np.ndarray) -> None: + h, w = frame.shape[:2] + os.makedirs(os.path.dirname(os.path.abspath(self.path)) or ".", exist_ok=True) + fourcc = cv2.VideoWriter_fourcc(*self.fourcc) if len(self.fourcc) == 4 else 0 + writer = cv2.VideoWriter(self.path, fourcc, self.fps, (w, h), True) + if not writer.isOpened(): + raise RuntimeError( + f"cv2.VideoWriter could not open {self.path!r} with fourcc={self.fourcc!r}" + ) + self._writer = writer + self._size = (w, h) + if self.verbose: + print(f"[recorder] {self.path} {w}x{h} @ {self.fps:g} fps fourcc={self.fourcc}") + + def _render(self, snap: Any) -> np.ndarray: + if snap.image is None or snap.flow is None: + raise RuntimeError( + "recording needs image and flow in the snapshot; build the pipeline " + 'with include=("image", "flow")' + ) + mask = self.valid_mask() if callable(self.valid_mask) else self.valid_mask + image = draw_arrows( + snap.image, + snap.flow, + threshold=self.threshold, + grid_spacing=self.grid_spacing, + arrow_scale=self.arrow_scale, + margin=self.margin, + valid_mask=mask, + ) + if self.draw_contact_flag and snap.contour is not None and snap.centroid is not None: + image = draw_contact(image, snap.contour, snap.centroid) + if self.overlay_values or self.label: + self._draw_hud(image, snap) + return image + + def _draw_hud(self, image: np.ndarray, snap: Any) -> None: + lines = [] + if self.label: + lines.append(str(self.label)) + if self.overlay_values: + lines.append(f"{snap.unit} fn={snap.normal_force:.3g} " + f"fx={snap.shear_x:.3g} fy={snap.shear_y:.3g}") + lines.append(f"{snap.fps:.0f} FPS {snap.age_ms:.1f} ms") + y = 14 + for text in lines: + cv2.putText(image, text, (4, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, + (0, 0, 0), 3, cv2.LINE_AA) + cv2.putText(image, text, (4, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, + (255, 255, 255), 1, cv2.LINE_AA) + y += 14 + + def _loop(self) -> None: + try: + while True: + try: + snap = self._queue.get(timeout=0.05) + except queue.Empty: + if self._stop.is_set(): + break + continue + if snap is None: + break + try: + frame = self._render(snap) + if self._writer is None: + self._open(frame) + self._writer.write(frame) + self._written += 1 + except BaseException as exc: + self._error = exc + break + finally: + if self._writer is not None: + self._writer.release() + self._writer = None + + # ------------------------------------------------------------------- close + def close(self, timeout: float = 30.0) -> dict: + """Drain, stop the writer thread and release the file.""" + if self._thread is not None: + self._stop.set() + # The stop flag is the real mechanism - the writer drains what is + # queued and exits on its next empty poll. The sentinel only makes + # that immediate, so it must never block: if the writer died on an + # error the queue may be full and nobody would ever drain it. + try: + self._queue.put_nowait(None) + except queue.Full: + pass + self._thread.join(timeout=timeout) + alive = self._thread.is_alive() + self._thread = None + else: + alive = False + stats = self.stats() + stats["thread_alive"] = alive + if self._error is not None: + stats["error"] = f"{type(self._error).__name__}: {self._error}" + if self.verbose: + print(f"[recorder] {self.path}: {stats['error']}") + return stats + + def stats(self) -> dict: + return { + "path": self.path, + "seen": self._seen, + "written": self._written, + "dropped": self._dropped, + "skipped": self._skipped, + "fps": self.fps, + "fourcc": self.fourcc, + } + + def __enter__(self) -> "StreamRecorder": + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + self.close() + return False + + +def mosaic_videos( + paths: Sequence[str], + out_path: str, + *, + cols: int | None = None, + labels: Iterable[str] | None = None, + scale: float = 1.0, + gap: int = 4, + fps: float | None = None, + fourcc: str = "mp4v", +) -> dict: + """Tile per-stream recordings into one grid video. + + Post-processing only: run it after the streams have finished, so it costs + nothing while they are running. Shorter inputs hold their last frame. + """ + paths = [str(p) for p in paths] + if not paths: + raise ValueError("mosaic_videos needs at least one input") + for path in paths: + if not os.path.isfile(path): + raise FileNotFoundError(path) + + caps = [cv2.VideoCapture(p) for p in paths] + try: + for path, cap in zip(paths, caps): + if not cap.isOpened(): + raise RuntimeError(f"cannot open {path!r}") + tiles = [] + for cap in caps: + ok, frame = cap.read() + if not ok: + raise RuntimeError("an input video has no frames") + tiles.append(frame) + + if scale != 1.0: + tiles = [cv2.resize(t, None, fx=scale, fy=scale, + interpolation=cv2.INTER_AREA) for t in tiles] + th, tw = tiles[0].shape[:2] + n = len(tiles) + cols = int(cols) if cols else int(np.ceil(np.sqrt(n))) + cols = max(1, min(cols, n)) + rows = int(np.ceil(n / cols)) + label_list = [str(x) for x in labels] if labels is not None else [] + + out_h = rows * th + (rows - 1) * gap + out_w = cols * tw + (cols - 1) * gap + canvas = np.zeros((out_h, out_w, 3), dtype=np.uint8) + positions = [] + for i in range(n): + r, c = divmod(i, cols) + positions.append((r * (th + gap), c * (tw + gap))) + # Frame 0 is already in hand; seed the canvas with it so the output is + # never short by one frame. + for i, tile in enumerate(tiles): + y, x = positions[i] + canvas[y:y + th, x:x + tw] = tile + + out_fps = float(fps) if fps else float(caps[0].get(cv2.CAP_PROP_FPS) or 30.0) + os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True) + writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*fourcc), out_fps, + (out_w, out_h), True) + if not writer.isOpened(): + raise RuntimeError(f"cannot open {out_path!r} for writing") + + def emit() -> None: + # Labels go on the canvas, not on the tiles: the read loop below + # overwrites each tile region every frame. + for i, text in enumerate(label_list): + if i >= n: + break + y, x = positions[i] + org = (x + 4, y + th - 6) + cv2.putText(canvas, text, org, cv2.FONT_HERSHEY_SIMPLEX, 0.45, + (0, 0, 0), 3, cv2.LINE_AA) + cv2.putText(canvas, text, org, cv2.FONT_HERSHEY_SIMPLEX, 0.45, + (255, 255, 255), 1, cv2.LINE_AA) + writer.write(canvas) + + written = 0 + try: + # One frame per input per iteration; a spent input keeps its last + # frame so the grid stays full until every input has run out. + while True: + emit() + written += 1 + progressed = False + for i, cap in enumerate(caps): + ok, frame = cap.read() + if not ok: + continue + progressed = True + if scale != 1.0: + frame = cv2.resize(frame, (tw, th), interpolation=cv2.INTER_AREA) + y, x = positions[i] + canvas[y:y + th, x:x + tw] = frame + if not progressed: + break + finally: + writer.release() + finally: + for cap in caps: + cap.release() + + return { + "path": out_path, + "frames": written, + "fps": out_fps, + "size": (out_w, out_h), + "grid": (rows, cols), + "scale": scale, + } diff --git a/src/orisys/pipeline.py b/src/orisys/pipeline.py new file mode 100644 index 0000000..11e1292 --- /dev/null +++ b/src/orisys/pipeline.py @@ -0,0 +1,465 @@ +"""Three-stage per-stream pipeline for the SDK's per-frame work. + +A single ``Sensor`` frame costs the sum of three stages:: + + produce -> capture + (stitch | UV unroll) + grey/resize (Sensor.get_img) + flow -> DIS optical flow + forward/backward + loop closure + hhd -> Helmholtz-Hodge decomposition + force readout (+ contact) + +Running them back to back in one thread bounds throughput by that *sum*. Running +them as three concurrent threads bounds it by the *slowest stage* only, which is +what makes multi-sensor deployments viable: measured on a Ryzen 9 9950X, five +finger streams go from 60.8 to 125.2 FPS per stream, and five box streams from +39.1 to 57.6 FPS per stream. + +Usage (finger):: + + with orisys.StreamPipeline.for_finger_runtime(runtime) as pipe: + while True: + snap = pipe.wait() + if snap is None: + break + print(snap.normal_force, snap.shear_x, snap.shear_y) + +Usage (box / mini):: + + with orisys.StreamPipeline(sensor, sensor.get_img) as pipe: + ... + +Data ownership +-------------- +The SDK reuses its working buffers between frames (``Sensor.img`` is +``_img_gray_buf``, ``Sensor.flow`` may be the flow-EMA buffer, ``Sensor.vnormal`` +is the HHD upsample buffer, and the finger remap output buffer is reused too), so +each stage must not hand a bare reference to the next one: + +* the producer copies the frame it publishes, because the next ``get_img`` call + overwrites the SDK buffer in place; +* the flow stage takes ownership of ``Sensor.img`` for the duration of the frame, + which is why the producer writes its own copy rather than letting the SDK + buffer be read while being rewritten; +* the flow stage also copies ``Sensor.flow`` before handing the frame on, because + ``Sensor.compute_hhd`` reads ``self.flow`` - and with the flow-EMA enabled that + attribute is a buffer the next frame rewrites. Handing the copy forward and + restoring ``Sensor.flow`` in the HHD stage is what keeps each frame's forces + paired with that frame's flow; +* the HHD stage owns every array that goes into a :class:`FrameSnapshot` (the + frame copy is already private), so a published snapshot is immutable and safe + to hand to a GUI thread. + +``Sensor.valid_mask`` is still rebuilt by the producer; it is geometry that does +not change between frames in practice, so the stages only ever read it. A sensor +whose mask changes every frame would need the SDK to expose a frame-scoped mask. + +Cost: one grey-frame copy (~0.1 MB) plus one flow copy (~1 MB at 365x365) per +frame, roughly 0.1 ms - about 0.02 cores of a 32-thread machine at five streams. + +Latency +------- +With the queues empty (a real camera at its own frame rate) the end-to-end +capture-to-force delay is the sum of the three stage times. ``depth`` bounds how +many frames may be in flight; a larger depth absorbs jitter at the cost of +latency when the producer outruns the consumer. +""" + +from __future__ import annotations + +import threading +import time +from dataclasses import dataclass, field +from queue import Empty, Full, Queue +from typing import Any, Callable, Iterable, Sequence + +import numpy as np + +__all__ = ["FrameSnapshot", "PipelineStats", "StreamPipeline"] + +#: Array fields a snapshot can carry. ``image`` is the working grey frame; the +#: rest come from ``Sensor.read_info`` (already copied by the SDK). +_ARRAY_FIELDS = ("image", "flow", "vnormal", "vshear", "depth", "centroid", "contour") + + +@dataclass(frozen=True) +class FrameSnapshot: + """An immutable view of one finished frame. + + Scalars are always populated; arrays are ``None`` unless requested through + ``StreamPipeline(include=...)``. + """ + + seq: int + timestamp: float + age_ms: float + fps: float + unit: str + normal_force: float + shear_x: float + shear_y: float + normal_force_raw: float + shear_x_raw: float + shear_y_raw: float + torque: float + image: np.ndarray | None = None + flow: np.ndarray | None = None + vnormal: np.ndarray | None = None + vshear: np.ndarray | None = None + depth: np.ndarray | None = None + centroid: np.ndarray | None = None + contour: np.ndarray | None = None + produce_ms: float = 0.0 + flow_ms: float = 0.0 + hhd_ms: float = 0.0 + + +@dataclass +class PipelineStats: + """Rolling diagnostics for one stream.""" + + produced: int = 0 + published: int = 0 + produce_ms: float = 0.0 + flow_ms: float = 0.0 + hhd_ms: float = 0.0 + age_ms: float = 0.0 + age_p95_ms: float = 0.0 + fps: float = 0.0 + _ages: list = field(default_factory=list, repr=False) + + def as_dict(self) -> dict: + return { + "produced": self.produced, + "published": self.published, + "fps": round(self.fps, 2), + "produce_ms": round(self.produce_ms, 3), + "flow_ms": round(self.flow_ms, 3), + "hhd_ms": round(self.hhd_ms, 3), + "age_ms": round(self.age_ms, 3), + "age_p95_ms": round(self.age_p95_ms, 3), + } + + +class StreamPipeline: + """Run ``produce -> flow -> hhd`` as three concurrent stages. + + ``produce`` is a zero-argument callable returning either the working frame + (``(H, W)`` grey or ``(H, W, 3)`` BGR) or ``None`` to signal end of input. + For a ``Sensor`` that is ``sensor.get_img``; for a ``FingerRuntime`` use + :meth:`for_finger_runtime`. + """ + + def __init__( + self, + sensor: Any, + produce: Callable[[], np.ndarray | None], + *, + contact: bool = False, + depth: int = 2, + include: Sequence[str] | Iterable[str] = ("image",), + on_frame: Callable[["FrameSnapshot"], None] | None = None, + verbose: bool = False, + fps_window: int = 30, + ) -> None: + unknown = set(include) - set(_ARRAY_FIELDS) + if unknown: + raise ValueError( + f"unknown include fields {sorted(unknown)}; expected {list(_ARRAY_FIELDS)}" + ) + self.sensor = sensor + self._produce = produce + self._contact = bool(contact) + self._depth = max(1, int(depth)) + self._include = tuple(dict.fromkeys(include)) + self._on_frame = on_frame + self._verbose = bool(verbose) + self._fps_window = max(2, int(fps_window)) + + self._q1: Queue = Queue(maxsize=self._depth) + self._q2: Queue = Queue(maxsize=self._depth) + self._done = threading.Event() + self._threads: list[threading.Thread] = [] + self._error: BaseException | None = None + self._running = False + + self._lock = threading.Lock() + self._latest: FrameSnapshot | None = None + self._published = 0 + self._wait_seq = 0 + self._stamps: list[float] = [] + self.stats = PipelineStats() + + # ------------------------------------------------------------------ ctors + @classmethod + def for_finger_runtime(cls, runtime: Any, **kwargs: Any) -> "StreamPipeline": + """Build a pipeline for a ``FingerRuntime`` (capture + UV unroll stage).""" + + def produce() -> np.ndarray | None: + ok, frame = runtime.grab_and_process(color=False) + return frame if ok else None + + return cls(runtime.tactile.raw_sensor, produce, **kwargs) + + # --------------------------------------------------------------- lifecycle + def start(self) -> "StreamPipeline": + if self._running: + return self + self._running = True + self._done.clear() + self._error = None + with self._lock: + self._wait_seq = self._published + self._threads = [ + threading.Thread(target=self._guard, args=(self._producer_loop,), + name="orisys-pipeline-produce", daemon=True), + threading.Thread(target=self._guard, args=(self._flow_loop,), + name="orisys-pipeline-flow", daemon=True), + threading.Thread(target=self._guard, args=(self._hhd_loop,), + name="orisys-pipeline-hhd", daemon=True), + ] + for t in self._threads: + t.start() + return self + + def stop(self, timeout: float = 5.0) -> None: + """Stop all stages and join them. Re-raises a stage exception, if any.""" + self._done.set() + self._drain(self._q1) + self._drain(self._q2) + for t in self._threads: + t.join(timeout=timeout) + self._threads = [] + self._running = False + if self._error is not None: + err, self._error = self._error, None + raise err + + def __enter__(self) -> "StreamPipeline": + return self.start() + + def __exit__(self, exc_type, exc, tb) -> bool: + try: + self.stop() + except BaseException: + if exc_type is None: + raise + return False + + # ------------------------------------------------------------------- state + @property + def running(self) -> bool: + return self._running and any(t.is_alive() for t in self._threads) + + @property + def error(self) -> BaseException | None: + return self._error + + def latest(self) -> FrameSnapshot | None: + """Most recent snapshot, or ``None`` if none was published yet.""" + with self._lock: + return self._latest + + def wait(self, timeout: float | None = None) -> FrameSnapshot | None: + """Block until a snapshot newer than the last returned one exists. + + This is a *latest-wins* view: if the stages outrun the caller, skipped + frames are gone, which is what a live GUI/control loop wants. Use + ``on_frame`` when every frame must be handled exactly once. + + Returns ``None`` on timeout, or once the pipeline has ended without a + newer frame. + """ + with self._lock: + seen = self._published + deadline = None if timeout is None else time.perf_counter() + timeout + while True: + with self._lock: + if self._published > self._wait_seq and self._latest is not None: + self._wait_seq = self._published + return self._latest + if self._error is not None: + err, self._error = self._error, None + raise err + if not self.running: + with self._lock: + if self._published > self._wait_seq and self._latest is not None: + self._wait_seq = self._published + return self._latest + return None + if deadline is not None and time.perf_counter() >= deadline: + return None + time.sleep(0.0005) + + # ----------------------------------------------------------------- internals + def _guard(self, target: Callable[[], None]) -> None: + try: + target() + except BaseException as exc: # noqa: BLE001 - surfaced by stop()/wait() + if self._error is None: + self._error = exc + self._done.set() + + def _drain(self, q: Queue) -> None: + while True: + try: + q.get_nowait() + except Empty: + return + + def _put(self, q: Queue, item: Any) -> bool: + while not self._done.is_set(): + try: + q.put(item, timeout=0.05) + return True + except Full: + continue + return False + + def _get(self, q: Queue) -> Any: + """Pop the next item, or ``None`` once the pipeline is stopping. + + Consumers must use this rather than ``q.get()``: when :meth:`stop` sets + the done flag the producers stop delivering sentinels, so a bare + blocking ``get()`` would wait forever and ``join()`` would time out. + """ + while not self._done.is_set(): + try: + return q.get(timeout=0.05) + except Empty: + continue + try: + return q.get_nowait() + except Empty: + return None + + def _producer_loop(self) -> None: + while not self._done.is_set(): + t0 = time.perf_counter() + frame = self._produce() + if frame is None: + break + t1 = time.perf_counter() + # Private copy: the SDK buffer is rewritten in place by the next + # produce() call while the flow stage is still reading this frame. + payload = (frame.copy(), t0, t1 - t0) + if not self._put(self._q1, payload): + break + self.stats.produced += 1 + self._put(self._q1, None) + + def _flow_loop(self) -> None: + while True: + item = self._get(self._q1) + if item is None: + break + frame, t_cap, produce_s = item + t0 = time.perf_counter() + # The frame is passed explicitly: the SDK's self.img slot is shared + # with the producer, so it must not be the transport. + self.sensor.compute_flow(frame) + flow_s = time.perf_counter() - t0 + flow = getattr(self.sensor, "flow", None) + # Private copy: with the flow EMA enabled Sensor.flow is a buffer the + # next frame rewrites. + flow_copy = flow.copy() if isinstance(flow, np.ndarray) else None + self._put(self._q2, (frame, flow_copy, t_cap, produce_s, flow_s)) + self._put(self._q2, None) + + def _hhd_loop(self) -> None: + while True: + item = self._get(self._q2) + if item is None: + break + frame, flow, t_cap, produce_s, flow_s = item + t0 = time.perf_counter() + # Explicit flow: compute_hhd reads self.flow, which the flow stage + # overwrites for the next frame, so it cannot be the transport. + self.sensor.compute_hhd(flow) + if self._contact: + self.sensor.compute_contact() + hhd_s = time.perf_counter() - t0 + snap = self._build(frame, flow, produce_s, flow_s, hhd_s, t_cap) + self._publish(snap) + if self._on_frame is not None: + # Runs on this stage's thread: a slow consumer exerts real + # backpressure instead of silently losing frames. + self._on_frame(snap) + if self._verbose: + print(f"[pipeline] finished after {self.stats.published} frames") + + # -------------------------------------------------------------- snapshot + def _copy_array(self, selector: Callable[[], Any]) -> np.ndarray | None: + try: + value = selector() + except Exception: + return None + if value is None or not isinstance(value, np.ndarray) or value.size == 0: + return None + return value.copy() + + def _build(self, frame: np.ndarray, flow: np.ndarray | None, produce_s: float, + flow_s: float, hhd_s: float, t_cap: float) -> FrameSnapshot: + sensor = self.sensor + info = sensor.info + include = self._include + now = time.perf_counter() + + def maybe(name: str, selector: Callable[[], Any]) -> np.ndarray | None: + return self._copy_array(selector) if name in include else None + + return FrameSnapshot( + seq=self.stats.published, + timestamp=time.time(), + age_ms=(now - t_cap) * 1000.0, + fps=self._fps(), + unit=str(getattr(sensor, "force_unit", "raw")), + normal_force=float(sensor.fnormal), + shear_x=float(sensor.fshearx), + shear_y=float(sensor.fsheary), + normal_force_raw=float(sensor._force_raw[0]), + shear_x_raw=float(sensor._force_raw[1]), + shear_y_raw=float(sensor._force_raw[2]), + torque=float(getattr(sensor, "ftor", 0.0)), + # Both already private per-frame copies owned by this snapshot. + image=frame if "image" in include else None, + flow=flow if "flow" in include else None, + # Written only by the HHD stage, so read_info is safe here. + vnormal=maybe("vnormal", lambda: sensor.read_info(info.VNORMAL)[0]), + vshear=maybe("vshear", lambda: sensor.read_info(info.VSHEAR)[0]), + depth=maybe("depth", lambda: sensor.read_info(info.DEPTH)[0]), + centroid=maybe("centroid", lambda: sensor.read_info(info.CENTROID)[0]), + contour=maybe("contour", lambda: sensor.read_info(info.CONTOUR)[0]), + produce_ms=produce_s * 1000.0, + flow_ms=flow_s * 1000.0, + hhd_ms=hhd_s * 1000.0, + ) + + def _fps(self) -> float: + now = time.perf_counter() + stamps = self._stamps + stamps.append(now) + if len(stamps) > self._fps_window: + del stamps[: len(stamps) - self._fps_window] + if len(stamps) < 2: + return 0.0 + span = stamps[-1] - stamps[0] + return (len(stamps) - 1) / span if span > 0 else 0.0 + + def _publish(self, snap: FrameSnapshot) -> None: + stats = self.stats + n = stats.published + 1 + stats.published = n + stats.produce_ms += (snap.produce_ms - stats.produce_ms) / n + stats.flow_ms += (snap.flow_ms - stats.flow_ms) / n + stats.hhd_ms += (snap.hhd_ms - stats.hhd_ms) / n + stats.age_ms += (snap.age_ms - stats.age_ms) / n + stats.fps = snap.fps + ages = stats._ages + ages.append(snap.age_ms) + if len(ages) > 200: + del ages[: len(ages) - 200] + # A percentile per published frame is measurable at 500+ FPS; refresh + # it periodically instead. + if n <= 10 or n % 10 == 0: + stats.age_p95_ms = float(np.percentile(ages, 95)) + with self._lock: + self._latest = snap + self._published = n diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index 0f57b5f..560bfee 100644 --- a/src/orisys/sensor.py +++ b/src/orisys/sensor.py @@ -689,11 +689,19 @@ class Sensor: self.compute_hhd() return True - def compute_flow(self): - if self.img is None: + def compute_flow(self, img=None): + """Compute the optical flow for one frame. + + ``img`` optionally supplies the frame to use instead of ``self.img``. + StreamPipeline passes it explicitly so a producer thread cannot swap + ``self.img`` between the hand-off and the tracker call. Positional + callers and existing code are unaffected. + """ + image = self.img if img is None else img + if image is None: raise Exception("Warning: Get image first") - flow = self.__tracker.track(self.img, valid_mask=self.valid_mask) + flow = self.__tracker.track(image, valid_mask=self.valid_mask) # flow = self.__tracker.track_raft(self.img) if hasattr(flow, 'download'): flow = flow.download() @@ -764,9 +772,16 @@ class Sensor: """``config``, ``measured`` or ``config (replay)``.""" return str(getattr(self, "_force_rate_source", "config")) - def compute_hhd(self): + def compute_hhd(self, flow=None): + """Decompose the optical flow and update the force readout. - flow = self.flow + ``flow`` optionally supplies the field to decompose instead of + ``self.flow``. StreamPipeline passes its own per-frame copy so the flow + stage cannot overwrite ``self.flow`` mid-decomposition. Existing + callers are unaffected. + """ + + flow = self.flow if flow is None else flow if flow is None: raise Exception("Warning: Compute flow first") diff --git a/src/orisys/util.py b/src/orisys/util.py index de2931a..908c61a 100644 --- a/src/orisys/util.py +++ b/src/orisys/util.py @@ -75,7 +75,10 @@ def draw_arrows( ) if mask is not None: valid &= mask[rr, cc] - valid &= mask[end_y, end_x] + # Arrow tips may fall outside the frame (the length cap is applied + # before the bounds check above), so clamp before indexing: those points + # are already False in ``valid``, and in-bounds indices are unchanged. + valid &= mask[np.clip(end_y, 0, h - 1), np.clip(end_x, 0, w - 1)] mag_range = mag.max() - mag.min() norm_mag = ( diff --git a/tests/test_draw_arrows.py b/tests/test_draw_arrows.py new file mode 100644 index 0000000..0cc74d8 --- /dev/null +++ b/tests/test_draw_arrows.py @@ -0,0 +1,116 @@ +"""orisys.util.draw_arrows: mask handling and the out-of-frame tip regression.""" + +from __future__ import annotations + +import cv2 +import numpy as np +import pytest + +from orisys.util import _prepare_draw_mask, draw_arrows + + +def flow_of(shape, dx, dy): + flow = np.zeros((*shape, 2), dtype=np.float32) + flow[..., 0] = dx / 2.0 # draw_arrows doubles the stored value + flow[..., 1] = dy / 2.0 + return flow + + +def background(frame): + return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR) + + +def test_out_of_frame_arrow_tip_is_dropped_not_raised(): + """Regression: mask[end_y, end_x] was indexed before the bounds check. + + A 40 px grid with the default edge_margin of 15 lets an arrow tip leave the + frame, which used to raise IndexError as soon as a valid_mask was supplied. + """ + frame = np.full((48, 48), 100, dtype=np.uint8) + flow = flow_of((48, 48), dx=40.0, dy=0.0) + mask = np.ones((48, 48), dtype=bool) + + out = draw_arrows(frame, flow, grid_spacing=40, margin=0.8, valid_mask=mask) + + assert out.shape == (48, 48, 3) + # The only grid point (20, 20) has a tip at x=58 in a 48 px frame, so the + # bounds check must have discarded it: nothing was drawn. + assert np.array_equal(out, background(frame)) + + +def test_out_of_frame_tip_matches_render_without_mask(): + """Clamping the mask lookup must not change which points are valid.""" + frame = np.full((48, 48), 100, dtype=np.uint8) + flow = flow_of((48, 48), dx=40.0, dy=0.0) + mask = np.ones((48, 48), dtype=bool) + + with_mask = draw_arrows(frame, flow, grid_spacing=40, valid_mask=mask) + without = draw_arrows(frame, flow, grid_spacing=40) + assert np.array_equal(with_mask, without) + + +def test_in_frame_arrows_are_drawn_through_an_all_true_mask(): + frame = np.full((96, 96), 90, dtype=np.uint8) + flow = flow_of((96, 96), dx=6.0, dy=3.0) + mask = np.ones((96, 96), dtype=bool) + + out = draw_arrows(frame, flow, grid_spacing=12, arrow_scale=1.0, + threshold=0.1, valid_mask=mask) + assert not np.array_equal(out, background(frame)) + + +def test_all_false_mask_draws_nothing(): + frame = np.full((96, 96), 90, dtype=np.uint8) + flow = flow_of((96, 96), dx=6.0, dy=3.0) + mask = np.zeros((96, 96), dtype=bool) + + out = draw_arrows(frame, flow, grid_spacing=12, arrow_scale=1.0, + threshold=0.1, valid_mask=mask) + assert np.array_equal(out, background(frame)) + + +def test_mask_can_exclude_one_side(): + """Endpoints outside the mask are skipped even when inside the frame.""" + frame = np.full((96, 96), 90, dtype=np.uint8) + flow = flow_of((96, 96), dx=8.0, dy=0.0) + mask = np.zeros((96, 96), dtype=bool) + mask[24:48, 24:72] = True # only the upper middle band + + out = draw_arrows(frame, flow, grid_spacing=12, arrow_scale=1.0, + threshold=0.1, valid_mask=mask) + assert not np.array_equal(out, background(frame)) + # Rows below the band stay untouched. + assert np.array_equal(out[80:, :, :], background(frame)[80:, :, :]) + + +def test_below_threshold_points_get_the_dot_colour(): + frame = np.full((96, 96), 90, dtype=np.uint8) + flow = flow_of((96, 96), dx=0.05, dy=0.05) + out = draw_arrows(frame, flow, grid_spacing=12, arrow_scale=1.0, + threshold=1.0, below_threshold_color=(0, 0, 255)) + # OpenCV colours are BGR, so (0, 0, 255) is red. + red = (out[:, :, 2] > 150) & (out[:, :, 0] < 100) & (out[:, :, 1] < 100) + assert red.sum() > 0 + + +def test_prepare_draw_mask_resizes_to_frame(): + mask = np.zeros((10, 20), dtype=bool) + mask[:, :10] = True + out = _prepare_draw_mask(mask, (40, 40)) + assert out.shape == (40, 40) + assert out[:, :20].all() + assert not out[:, 20:].any() + + +def test_prepare_draw_mask_passes_none_through(): + assert _prepare_draw_mask(None, (10, 10)) is None + + +def test_non_square_grid_does_not_raise_with_mask(): + """Height and width differ, so rr/cc broadcasting must stay (N, N).""" + frame = np.full((60, 120), 80, dtype=np.uint8) + flow = flow_of((60, 120), dx=5.0, dy=5.0) + out = draw_arrows(frame, flow, grid_spacing=10, arrow_scale=1.0, + threshold=0.1, valid_mask=np.ones((60, 120), dtype=bool)) + assert out.shape == (60, 120, 3) + assert not np.array_equal(out, background(frame)) diff --git a/tests/test_export.py b/tests/test_export.py new file mode 100644 index 0000000..900ae1c --- /dev/null +++ b/tests/test_export.py @@ -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")) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..ddcf4ec --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,329 @@ +"""StreamPipeline: staging, snapshot immutability, error propagation, shutdown.""" + +from __future__ import annotations + +import time + +import numpy as np +import pytest + +from orisys import FrameSnapshot, StreamPipeline + + +class FakeSensor: + """Minimal Sensor stand-in with the same mutable-buffer behaviour. + + It deliberately reuses one image buffer (like ``Sensor._img_gray_buf``) so the + pipeline's private-copy contract is actually exercised. + """ + + class info: + VRAW = "vraw" + VNORMAL = "vnormal" + VSHEAR = "vshear" + DEPTH = "depth" + CENTROID = "centroid" + CONTOUR = "contour" + + def __init__(self, *, fail_at: int | None = None, flow_delay: float = 0.0, + hhd_delay: float = 0.0): + self.img = None + self.flow = None + self.vnormal = None + self.vshear = None + self.depth_map = None + self.centroid = (0, 0) + self.contour = np.empty(0) + self.ftor = 0.0 + self._force_raw = (0.0, 0.0, 0.0) + self.fnormal = self.fshearx = self.fsheary = 0.0 + self.force_unit = "raw" + self._img_buf = np.zeros((8, 8), dtype=np.uint8) + self._flow_buf = np.zeros((8, 8, 2), dtype=np.float32) + self._flow_delay = flow_delay + self._hhd_delay = hhd_delay + self._fail_at = fail_at + self._flow_calls = 0 + self._seen_images: list[np.ndarray] = [] + self._seen_flow_values: list[float] = [] + self.contact_calls = 0 + + def make_frame(self, seq: int) -> np.ndarray: + # writes INTO the shared buffer, exactly like cv2.resize(dst=...) + self._img_buf.fill(seq % 251) + self.img = self._img_buf + return self._img_buf + + def compute_flow(self, img=None, **kwargs) -> None: + image = self.img if img is None else img + self._flow_calls += 1 + if self._fail_at is not None and self._flow_calls == self._fail_at: + raise RuntimeError("flow exploded") + self._seen_images.append(np.array(image, copy=True)) + time.sleep(self._flow_delay) + # rewrite the shared buffer, like the real DIS/EMA path does + self._flow_buf.fill(float(self._flow_calls)) + self.flow = self._flow_buf + + def compute_hhd(self, flow=None, **kwargs) -> None: + # Mirrors Sensor.compute_hhd, which derives the forces from the field it + # is given - so this only yields frame N's forces if the pipeline passes + # frame N's flow explicitly. + time.sleep(self._hhd_delay) + field = self.flow if flow is None else flow + value = (float(field[0, 0, 0]) + if isinstance(field, np.ndarray) and field.size else 0.0) + self._seen_flow_values.append(value) + self.fnormal = value * 10.0 + self.fshearx = value + self.fsheary = -value + self._force_raw = (self.fnormal, self.fshearx, self.fsheary) + self.ftor = value + self.vnormal = self._flow_buf + self.vshear = self._flow_buf + + def compute_contact(self) -> None: + self.contact_calls += 1 + self.centroid = (1, 2) + self.contour = np.zeros((1, 1, 2), dtype=np.int32) + self.depth_map = np.zeros((8, 8), dtype=np.float32) + + def read_info(self, *args): + source = { + self.info.VRAW: lambda: self.flow.copy(), + self.info.VNORMAL: lambda: self.vnormal.copy(), + self.info.VSHEAR: lambda: self.vshear.copy(), + self.info.DEPTH: lambda: self.depth_map.copy(), + self.info.CENTROID: lambda: np.array(self.centroid), + self.info.CONTOUR: lambda: self.contour.copy(), + } + return tuple(source[a]() for a in args) + + +def make_producer(sensor: FakeSensor, count: int): + state = {"i": 0} + + def produce(): + if state["i"] >= count: + return None + frame = sensor.make_frame(state["i"]) + state["i"] += 1 + return frame + + return produce + + +def run_all(sensor: FakeSensor, count: int, **kwargs): + """Drive the pipeline with an exhaustive consumer: every frame, in order.""" + frames: list = [] + kwargs.setdefault("depth", 2) + pipe = StreamPipeline(sensor, make_producer(sensor, count), + on_frame=frames.append, **kwargs) + pipe.start() + deadline = time.perf_counter() + 10.0 + while len(frames) < count and time.perf_counter() < deadline: + if pipe.error is not None: + break + time.sleep(0.001) + pipe.stop() + return pipe, frames + + +def test_on_frame_receives_every_frame_in_order(): + sensor = FakeSensor() + pipe, frames = run_all(sensor, 12) + assert len(frames) == 12 + assert [s.seq for s in frames] == list(range(12)) + assert pipe.stats.produced == 12 + assert pipe.stats.published == 12 + + +def test_forces_come_from_the_matching_frame(): + """Each snapshot must carry the forces computed from its own frame's flow.""" + sensor = FakeSensor() + _, frames = run_all(sensor, 10) + # frame i's flow value is (i + 1); the fake derives the force from self.flow, + # so a mismatched pairing would show up as a repeated or skipped value. + assert sensor._seen_flow_values == [float(i + 1) for i in range(10)] + assert [s.normal_force for s in frames] == [10.0 * (i + 1) for i in range(10)] + assert [s.shear_x for s in frames] == [float(i + 1) for i in range(10)] + + +def test_stages_do_not_read_a_rewritten_image_buffer(): + """The producer reuses one buffer; no stage may observe a later frame.""" + sensor = FakeSensor(flow_delay=0.004) + run_all(sensor, 8) + seen = sensor._seen_images + assert len(seen) == 8 + for i, arr in enumerate(seen): + assert np.all(arr == i % 251), f"frame {i} was overwritten before flow ran" + + +def test_snapshot_arrays_are_private_copies(): + sensor = FakeSensor() + _, frames = run_all(sensor, 4, include=("image", "flow", "vnormal")) + first = frames[0] + assert isinstance(first, FrameSnapshot) + assert first.image is not None and first.flow is not None + saved = first.flow.copy() + # the SDK keeps rewriting its buffers; the snapshot must not change + sensor._flow_buf.fill(999.0) + sensor._img_buf.fill(7) + assert np.array_equal(first.flow, saved) + assert not np.any(first.flow == 999.0) + + +def test_include_filters_which_arrays_are_copied(): + sensor = FakeSensor() + _, frames = run_all(sensor, 3, include=("vnormal",)) + assert frames[0].vnormal is not None + assert frames[0].image is None and frames[0].flow is None + assert frames[0].depth is None and frames[0].contour is None + + +def test_rejects_unknown_include_field(): + sensor = FakeSensor() + with pytest.raises(ValueError): + StreamPipeline(sensor, make_producer(sensor, 1), include=("floww",)) + + +def test_contact_stage_runs_when_requested(): + sensor = FakeSensor() + _, frames = run_all(sensor, 5, contact=True, + include=("depth", "centroid", "contour")) + assert sensor.contact_calls == 5 + assert frames[-1].depth is not None + assert frames[-1].centroid is not None + assert frames[-1].contour is not None + + +def test_stage_exception_surfaces_on_stop(): + sensor = FakeSensor(fail_at=3) + pipe = StreamPipeline(sensor, make_producer(sensor, 20), depth=2) + pipe.start() + time.sleep(0.4) + with pytest.raises(RuntimeError, match="flow exploded"): + pipe.stop() + + +def test_stage_exception_surfaces_on_wait(): + sensor = FakeSensor(fail_at=2) + pipe = StreamPipeline(sensor, make_producer(sensor, 20), depth=2) + pipe.start() + with pytest.raises(RuntimeError): + for _ in range(50): + pipe.wait(timeout=0.1) + pipe.stop() + + +def test_stop_is_idempotent(): + sensor = FakeSensor() + pipe, _ = run_all(sensor, 4) + pipe.stop() # second stop must not raise + assert not pipe.running + assert pipe.latest() is not None + + +def test_stop_while_running_does_not_hang(): + """Stopping mid-stream must unblock the consumer threads. + + The producer stops delivering sentinels once the done flag is set, so a + bare blocking Queue.get() in a consumer would wait forever. + """ + sensor = FakeSensor(flow_delay=0.002) + pipe = StreamPipeline(sensor, make_producer(sensor, 100000), depth=2) + pipe.start() + assert pipe.wait(timeout=2.0) is not None + t0 = time.perf_counter() + pipe.stop(timeout=5.0) + elapsed = time.perf_counter() - t0 + assert elapsed < 2.0, f"stop() took {elapsed:.2f}s; a stage was still blocked" + assert not pipe.running + + +def test_context_manager_stops_threads(): + sensor = FakeSensor() + frames: list = [] + with StreamPipeline(sensor, make_producer(sensor, 6), + on_frame=frames.append) as pipe: + deadline = time.perf_counter() + 5.0 + while len(frames) < 6 and time.perf_counter() < deadline: + time.sleep(0.001) + assert len(frames) == 6 + assert not pipe.running + + +def test_wait_exposes_a_latest_wins_view(): + sensor = FakeSensor() + pipe = StreamPipeline(sensor, make_producer(sensor, 5)) + pipe.start() + last = None + while True: + snap = pipe.wait(timeout=2.0) + if snap is None: + break + if last is not None: + assert snap.seq > last.seq + last = snap + # the pipeline has drained; the final snapshot is the last frame + assert last is not None and last.normal_force == 50.0 + assert pipe.latest() is last + # nothing new is coming, so a further poll returns immediately + assert pipe.wait(timeout=0.2) is None + pipe.stop() + + +def test_wait_timeout_returns_none_without_new_frames(): + sensor = FakeSensor() + pipe = StreamPipeline(sensor, make_producer(sensor, 1)) + pipe.start() + assert pipe.wait(timeout=1.0) is not None + t0 = time.perf_counter() + assert pipe.wait(timeout=0.1) is None + assert time.perf_counter() - t0 < 1.0 + pipe.stop() + + +def test_stats_track_stage_times_and_age(): + sensor = FakeSensor(flow_delay=0.002, hhd_delay=0.001) + pipe, _ = run_all(sensor, 15) + d = pipe.stats.as_dict() + assert d["published"] == 15 + assert d["flow_ms"] >= 1.5 # the injected delay shows up + assert d["hhd_ms"] >= 0.8 + assert d["age_ms"] > 0.0 + assert d["fps"] > 0.0 + assert d["produce_ms"] >= 0.0 + + +def test_depth_does_not_change_results(): + for depth in (1, 2, 4): + sensor = FakeSensor() + _, frames = run_all(sensor, 10, depth=depth) + assert [s.normal_force for s in frames] == [10.0 * (i + 1) for i in range(10)] + + +def test_for_finger_runtime_maps_eof_to_none(): + """``for_finger_runtime`` drives ``grab_and_process`` and maps EOF to None.""" + sensor = FakeSensor() + state = {"i": 0} + + class FakeRuntime: + tactile = type("T", (), {"raw_sensor": sensor})() + + def grab_and_process(self, color=False): + if state["i"] >= 3: + return False, None + frame = sensor.make_frame(state["i"]) + state["i"] += 1 + return True, frame + + frames: list = [] + pipe = StreamPipeline.for_finger_runtime(FakeRuntime(), on_frame=frames.append) + pipe.start() + deadline = time.perf_counter() + 5.0 + while len(frames) < 3 and time.perf_counter() < deadline: + time.sleep(0.001) + pipe.stop() + assert len(frames) == 3 + assert pipe.stats.produced == 3