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