"""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