added 5 sensors multi-stream support, 5 finger reach > 100 fps stable using .mp4 video
This commit is contained in:
ykcheungab
@@ -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,
|
||||
}
|
||||
Reference in New Issue
Block a user