Files
orisys-sdk/src/orisys/pipeline.py
T
2026-09-30 21:23:08 +08:00

675 lines
27 KiB
Python

"""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 collections.abc import Callable, Iterable, Sequence
from dataclasses import dataclass, field
from enum import Enum
from queue import Empty, Full, Queue
from typing import Any
import numpy as np
__all__ = ["FrameSnapshot", "PipelineStats", "PauseLevel", "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")
#: How many periods the stages wait on the pause event before re-checking for a
#: stop. Small enough to resume promptly, large enough not to spin a core.
_PAUSE_POLL_S = 0.05
class PauseLevel(str, Enum):
"""What a pause holds still.
``BOTH`` stops the pipeline read and the compute stages: nothing is captured
and nothing is computed, so the sensor reports nothing until it resumes.
``CAPTURE_ONLY`` keeps reading frames and skips the flow and HHD stages. The
read is still issued, so a live camera keeps streaming: its exposure,
white balance and frame rate stay where the operator set them, and resuming
does not have to re-converge. What stops is the work - the three-stage
pipeline drops to one stage of cost, which is what makes a long hold cheap.
Both levels rebase the tracker on resume, so the hold does not turn into a
phantom deformation reading.
"""
BOTH = "both"
CAPTURE_ONLY = "capture_only"
@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
#: Capture resolution before any stitch/unroll/resize, as ``(width, height)``, or ``None``
#: when the producer did not report one. Read from ``Sensor.img_size_raw`` because callers
#: compare it with ``image.shape`` to see how much the frame was reduced.
image_size_raw: tuple[int, int] | None = None
@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
#: Set by :meth:`reset` and consumed by the flow stage, so that resetting sensor state
#: never races the stage that owns that state.
self._reset_pending = threading.Event()
#: Pause state. ``_pause_level`` is what a pause would hold when armed;
#: ``_currently_capturing`` records whether the producer is still
#: reading, which decides whether a resume has to rebase the tracker.
self._pause = threading.Event()
self._pause_level: PauseLevel = PauseLevel.BOTH
self._currently_capturing = True
self._sensor_source: Any = None
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 reset(self) -> None:
"""Reset the sensor's tracker and force-filter state, safely.
Call this from the consumer thread (after a key press, on a drift alarm). The reset is
queued and applied by the flow stage before its next frame, never inline: ``Sensor.reset``
rebases ``initial_img`` on the current ``Sensor.img``, and the flow stage is reading that
buffer concurrently, so resetting from the caller's thread would capture a half-written
reference frame.
Idempotent: a second call before the next frame does not queue twice.
"""
self._reset_pending.set()
@property
def paused(self) -> bool:
return self._pause.is_set()
@property
def pause_level(self) -> PauseLevel:
"""The level in effect while paused (the last level requested)."""
return self._pause_level
def pause(self, level: PauseLevel | str = PauseLevel.CAPTURE_ONLY) -> None:
"""Hold the pipeline still until :meth:`resume`.
``level`` selects what stays running:
* :attr:`PauseLevel.BOTH` - stop reading *and* computing. The sensor
reports nothing at all while held.
* :attr:`PauseLevel.CAPTURE_ONLY` (default) - keep reading so a live
camera stays in its converged state (exposure, white balance, frame
rate), but skip the flow and HHD stages. Cost drops to the read alone.
No snapshots are published while paused, so :meth:`wait` returns
``None``. The tracker is rebased on resume, because a held reference
frame would otherwise report the hold itself as a large deformation -
use ``reset=False`` to resume without that rebase.
Safe to call from the consumer thread, and idempotent; re-pausing with a
different level takes effect immediately.
"""
if not isinstance(level, PauseLevel):
level = PauseLevel(level)
if self._done.is_set():
return
self._set_reading(level is PauseLevel.CAPTURE_ONLY)
self._pause_level = level
self._pause.set()
def resume(self, *, reset: bool = True) -> None:
"""Release a pause and continue processing.
With ``reset`` (the default) the tracker is rebased on the next frame,
which is what keeps the hold out of the force reading. Pass
``reset=False`` when the caller wants the pre-pause reference frame kept
- for example when the held period is known to be motion-free.
"""
if not self._pause.is_set():
return
if reset and self._currently_capturing:
# Frames read during the pause were never computed, so the sensor's
# reference frame is stale by however long the hold lasted. Ask for a
# rebase; the flow stage applies it before its next frame.
self._reset_pending.set()
# Arm capture *before* lifting the hold. The producer stops as soon as
# the pause flag is set and only returns once it is clear, so a level-1
# resume must have `_currently_capturing` true by then or the producer
# settles with capture mode off.
self._currently_capturing = True
self._pause.clear()
self._set_reading(True)
def _set_reading(self, reading: bool) -> None:
"""Turn the producer's capture on or off (level 1 versus level 2)."""
self._currently_capturing = bool(reading)
source = self._resolve_sensor_source()
if source is None:
return
try:
if reading:
source.start_background()
else:
source.stop_background()
except Exception: # noqa: BLE001 - a video file may have no such control
if self._verbose:
print("[pipeline] pause: could not toggle background capture", flush=True)
def _resolve_sensor_source(self) -> Any:
"""The sensor's capture source, or ``None`` for file/derived sources."""
if self._sensor_source is not None:
return self._sensor_source
source = getattr(self.sensor, "_capture_source", None)
if source is None:
return None
# File playback cannot be held this way, and stopping a fake source only
# breaks the read, so leave those alone.
try:
if not hasattr(source, "start_background") or not source.is_open():
return None
except Exception: # noqa: BLE001
return None
self._sensor_source = source
return source
def _skip_while_paused(self) -> bool:
"""True while a pause discards this stage's work."""
return self._pause.is_set() and not self._wait_out_pause()
def _wait_out_pause(self) -> bool:
"""Block while paused; ``False`` means the pipeline is stopping."""
while self._pause.is_set() and not self._done.is_set():
self._pause.wait(_PAUSE_POLL_S)
return not self._done.is_set()
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, while the pipeline is paused, or once the
pipeline has ended without a newer frame. ``running`` distinguishes a
hold from the end of input.
"""
if self._pause.is_set():
# A hold publishes nothing. Returning the pre-pause snapshot here
# would look like live data to a control loop.
return None
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 _keep_camera_fed(self) -> None:
"""Read and discard one frame, without touching the shared sensor state.
``get_img`` would write ``Sensor.img`` while the flow stage may still be
draining, so the capture source is read directly. The read is what keeps
a live camera's exposure and auto-white-balance converged for however
long the pause lasts; nothing downstream sees the frame.
"""
source = self._sensor_source
if source is None:
time.sleep(_PAUSE_POLL_S)
return
try:
source.read_latest(copy=False, wait_s=0.0)
except Exception: # noqa: BLE001 - a dropped frame is not a failure
time.sleep(_PAUSE_POLL_S)
def _producer_loop(self) -> None:
while not self._done.is_set():
if self._pause.is_set() and not self._currently_capturing:
# Level 1 hold: capture itself is stopped, so there is nothing to
# read. Wait it out, then drop the stale backlog.
if not self._wait_out_pause():
break
self._drain(self._q1)
continue
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
if self._pause.is_set():
# Level 2 holds the frame in camera only: skip the flow work and
# let the queue drop this frame. A motion-free hold adds nothing
# to the track, and the reference frame is rebased on resume.
if not self._wait_out_pause():
break
self._drain(self._q1)
self._drain(self._q2)
continue
# Applied here, on the thread that owns the sensor's flow state, before this frame is
# computed: see reset() for why this must not run on the caller's thread.
if self._reset_pending.is_set():
self._reset_pending.clear()
self.sensor.reset()
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
if self._pause.is_set():
# Discard rather than compute: this is the stage the level-2 hold
# is meant to avoid paying for.
self._drain(self._q2)
if not self._wait_out_pause():
break
continue
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
raw_size = getattr(sensor, "img_size_raw", None)
if isinstance(raw_size, (tuple, list)) and len(raw_size) == 2:
image_size_raw: tuple[int, int] | None = (int(raw_size[0]), int(raw_size[1]))
else:
image_size_raw = 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,
image_size_raw=image_size_raw,
)
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