ykcheungab
366 lines
14 KiB
Python
366 lines
14 KiB
Python
"""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,
|
|
}
|