orisys-sdk/examples/pipeline/service.py

"""Finger tactile acquisition loop for example dashboards."""

from __future__ import annotations

import time
from dataclasses import dataclass

import numpy as np
from orisys.finger import FingerRuntime

from viz import FrameRecorder, build_uv_surface_mesh, resolve_viz_pose_path

from .config import PipelineConfig
from .frame import FingerFrameView
from .profiling import StageProfiler, run_deformation_timed


@dataclass
class TickTimings:
    capture_ms: float = 0.0
    uv_unroll_ms: float = 0.0
    flow_ms: float = 0.0
    hhd_ms: float = 0.0
    contact_ms: float = 0.0
    read_info_ms: float = 0.0
    pipeline_ms: float = 0.0


class FingerPipeline:
    """Owns FingerRuntime and produces FingerFrameView snapshots per tick."""

    def __init__(self, config: PipelineConfig):
        self.config = config
        self.runtime: FingerRuntime | None = None
        self.recorder = FrameRecorder()
        self._start_time: float | None = None
        self.uv_mesh = None
        self.world_to_cam: np.ndarray | None = None
        self.mesh_lookat: np.ndarray | None = None
        self.profiler: StageProfiler | None = (
            StageProfiler(report_every=config.profile_every) if config.profile else None
        )

    @property
    def is_open(self) -> bool:
        return self.runtime is not None

    @property
    def tactile(self):
        if self.runtime is None:
            raise RuntimeError("FingerPipeline is not open")
        return self.runtime.tactile

    def open(self) -> None:
        if self.runtime is not None:
            return
        self.runtime = FingerRuntime.open(
            config=self.config.config,
            calib=self.config.calib,
            video=self.config.video,
            verbose=self.config.verbose,
        )
        self._prepare_mesh_assets()
        self.runtime.warmup(color=self.config.use_color)
        self._start_time = time.time()

    def _prepare_mesh_assets(self) -> None:
        if self.runtime is None:
            return
        domain = self.runtime.domain
        if domain.surface is None or getattr(domain.surface, "world_xyz", None) is None:
            return
        self.uv_mesh = build_uv_surface_mesh(domain.surface.world_xyz, domain.valid_mask)
        if self.config.viz_pose.lower() != "none":
            self.world_to_cam = np.load(resolve_viz_pose_path(self.config.viz_pose))
            self.mesh_lookat = self.uv_mesh.rest_vertices.mean(axis=0)

    def close(self) -> None:
        if self.profiler is not None and self.profiler.frame_count:
            self.profiler.print_report(prefix="[profile final]")
        if self.runtime is not None:
            self.runtime.close()
            self.runtime = None
        self._start_time = None

    def reset(self) -> None:
        if self.runtime is not None:
            self.runtime.tactile.reset()

    def tick(self) -> FingerFrameView | None:
        frame, _timings = self.tick_timed()
        return frame

    def tick_timed(self) -> tuple[FingerFrameView | None, TickTimings]:
        timings = TickTimings()
        if self.runtime is None or self._start_time is None:
            return None, timings

        pipeline_t0 = time.perf_counter()
        profile = self.profiler is not None

        if profile:
            t0 = time.perf_counter()
            ok, frame = self.runtime.read_source()
            timings.capture_ms = (time.perf_counter() - t0) * 1000.0
            if not ok:
                return None, timings
            t0 = time.perf_counter()
            unrolled = self.runtime.process_frame(frame, color=self.config.use_color)
            timings.uv_unroll_ms = (time.perf_counter() - t0) * 1000.0
        else:
            ok, unrolled = self.runtime.grab_and_process(color=self.config.use_color)
            if not ok:
                return None, timings

        tactile = self.runtime.tactile
        self.recorder.write_if_recording(tactile.frame)

        if profile:
            flow_ms, hhd_ms = run_deformation_timed(
                tactile, check_motion=True, threshold=0
            )
            timings.flow_ms = flow_ms
            timings.hhd_ms = hhd_ms
            self.profiler.record("capture", timings.capture_ms)
            self.profiler.record("uv_unroll", timings.uv_unroll_ms)
            self.profiler.record("flow", timings.flow_ms)
            self.profiler.record("hhd", timings.hhd_ms)
        else:
            tactile.compute_deformation(decompose=True, check_motion=True, threshold=0)

        t0 = time.perf_counter()
        tactile.compute_contact()
        timings.contact_ms = (time.perf_counter() - t0) * 1000.0
        if profile:
            self.profiler.record("contact", timings.contact_ms)

        t0 = time.perf_counter()
        view = FingerFrameView.from_tactile(
            tactile,
            unrolled,
            timestamp=time.time() - self._start_time,
        )
        timings.read_info_ms = (time.perf_counter() - t0) * 1000.0
        if profile:
            self.profiler.record("read_info", timings.read_info_ms)

        timings.pipeline_ms = (time.perf_counter() - pipeline_t0) * 1000.0
        if profile:
            self.profiler.record("pipeline", timings.pipeline_ms)

        return view, timings

    def record_total_frame(self, total_ms: float) -> None:
        if self.profiler is None:
            return
        self.profiler.record("total", total_ms)
        self.profiler.tick()

    def start_recording(self) -> bool:
        if self.runtime is None:
            return False
        return self.recorder.start(self.runtime.tactile.frame)

    def stop_recording(self) -> None:
        self.recorder.stop()