split example visualizaiton pipeline

2026-06-18 13:47:16 +08:00
parent fbee1c37ed
commit f19f092b64
42 changed files with 3067 additions and 702 deletions
@@ -0,0 +1,328 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import cv2
+import numpy as np
+
+from ..configs import load_config
+from ..frame_processors import UVUnrollFrameProcessor
+from ..sensor import Sensor
+from ..unroll import (
+    build_or_load_domain,
+    load_calib_json,
+    resolve_uv_runtime_sizes,
+    scale_camera_geometry,
+)
+from .exports import (
+    DEFAULT_FINGER_EXPORT_DIR,
+    DEFAULT_FINGER_MAP_CACHE,
+    FingerExportPaths,
+    load_finger_export_bundle,
+)
+
+
+def normalize_video_source(video):
+    """Accept camera index, numeric string, or filesystem video path."""
+    if video is None:
+        return 0
+    if isinstance(video, str):
+        stripped = video.strip().strip("/\\")
+        if not stripped:
+            return 0
+        if stripped.isdigit():
+            return int(stripped)
+        return stripped
+    return video
+
+
+def open_source(video, config):
+    capture = cv2.VideoCapture(normalize_video_source(video))
+    capture.set(cv2.CAP_PROP_FRAME_WIDTH, config.get("camera", {}).get("resolution_raw")[0])
+    capture.set(cv2.CAP_PROP_FRAME_HEIGHT, config.get("camera", {}).get("resolution_raw")[1])
+    capture.set(cv2.CAP_PROP_FPS, config.get("camera", {}).get("fps", 30))
+    capture.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M', 'J', 'P', 'G'))
+    return capture
+
+
+class TactileProcessor:
+    """Narrow tactile-analysis wrapper around ``Sensor`` external-frame mode."""
+
+    def __init__(self, sensor: Sensor, *, default_valid_mask=None):
+        self._sensor = sensor
+        self._default_valid_mask = default_valid_mask
+
+    @property
+    def raw_sensor(self) -> Sensor:
+        """Escape hatch for legacy callers that still need the full ``Sensor``."""
+        return self._sensor
+
+    @property
+    def info(self):
+        return self._sensor.info
+
+    @property
+    def frame(self):
+        return self._sensor.frame
+
+    @property
+    def img(self):
+        return self._sensor.img
+
+    @property
+    def flow(self):
+        return self._sensor.flow
+
+    @property
+    def vnormal(self):
+        return self._sensor.vnormal
+
+    @property
+    def vshear(self):
+        return self._sensor.vshear
+
+    @property
+    def depth_map(self):
+        return self._sensor.depth_map
+
+    @property
+    def mask_stats(self):
+        return self._sensor.mask_stats
+
+    @property
+    def valid_mask(self):
+        return self._sensor.valid_mask
+
+    def push_frame(self, frame, *, already_gray=False, valid_mask=None):
+        """Feed one processed tactile frame into the analysis pipeline."""
+        mask = self._default_valid_mask if valid_mask is None else valid_mask
+        return self._sensor.get_img(
+            frame=frame,
+            already_processed=True,
+            already_gray=already_gray,
+            valid_mask=mask,
+        )
+
+    def warmup(self):
+        return self._sensor.warmup_deformation()
+
+    def compute_deformation(self, *args, **kwargs):
+        return self._sensor.compute_deformation(*args, **kwargs)
+
+    def compute_contact(self, *args, **kwargs):
+        return self._sensor.compute_contact(*args, **kwargs)
+
+    def read_info(self, *args):
+        return self._sensor.read_info(*args)
+
+    def reset(self):
+        return self._sensor.reset()
+
+    def disconnect(self):
+        return self._sensor.disconnect()
+
+
+@dataclass(frozen=True)
+class FingerGeometry:
+    """Resolved one-time geometry/bootstrap state for a finger UV runtime."""
+
+    config: dict
+    camera_config: dict
+    export_paths: FingerExportPaths
+    data: dict
+    source_K: np.ndarray
+    dist_coeffs: np.ndarray | None
+    native_width: int
+    native_height: int
+    source_width: int
+    source_height: int
+    out_width: int
+    out_height: int
+    source_scale: float
+    map_cache: str
+    edge_expand: float
+
+
+def resolve_finger_geometry(
+    *,
+    export_dir,
+    config="./config/dd02-ov.json",
+    map_cache=DEFAULT_FINGER_MAP_CACHE,
+    source_scale=None,
+    calib=None,
+) -> FingerGeometry:
+    """Resolve config/export/camera geometry for a UV finger runtime."""
+    config = load_config(config)
+    camera_config = config.get("camera", {})
+    out_width, out_height, source_scale = resolve_uv_runtime_sizes(
+        camera_config,
+        source_scale=source_scale,
+    )
+    bundle = load_finger_export_bundle(
+        export_dir,
+        map_cache=map_cache,
+        source_scale=source_scale,
+    )
+
+    if calib:
+        source_K, dist_coeffs, img_size = load_calib_json(calib)
+        native_width, native_height = img_size
+    else:
+        source_K = bundle.data["K"]
+        dist_coeffs = None
+        native_width = int(bundle.data["width"])
+        native_height = int(bundle.data["height"])
+
+    source_K, source_width, source_height = scale_camera_geometry(
+        source_K,
+        native_width,
+        native_height,
+        source_scale,
+    )
+
+    return FingerGeometry(
+        config=config,
+        camera_config=camera_config,
+        export_paths=bundle.export_paths,
+        data=bundle.data,
+        source_K=source_K,
+        dist_coeffs=dist_coeffs,
+        native_width=native_width,
+        native_height=native_height,
+        source_width=source_width,
+        source_height=source_height,
+        out_width=out_width,
+        out_height=out_height,
+        source_scale=float(source_scale),
+        map_cache=bundle.map_cache,
+        edge_expand=bundle.edge_expand,
+    )
+
+
+class FingerRuntime:
+    """Reusable UV finger runtime composed from capture, UV remap, and tactile analysis."""
+
+    def __init__(
+        self,
+        *,
+        geometry: FingerGeometry,
+        domain,
+        tactile: TactileProcessor,
+        frame_processor: UVUnrollFrameProcessor,
+        capture,
+    ):
+        self.geometry = geometry
+        self.domain = domain
+        self.tactile = tactile
+        self.sensor = tactile
+        self.frame_processor = frame_processor
+        self.capture = capture
+
+    @classmethod
+    def open(
+        cls,
+        *,
+        export_dir=DEFAULT_FINGER_EXPORT_DIR,
+        config="./config/dd02-ov.json",
+        map_cache=DEFAULT_FINGER_MAP_CACHE,
+        source_scale=None,
+        calib=None,
+        video=0,
+        camera_id=None,
+        flip_source_vertical=True,
+        verbose=False,
+        mirror=False,
+    ) -> "FingerRuntime":
+        """Open a live finger runtime from export bundle, config, and video source."""
+        if camera_id is not None:
+            video = camera_id
+        video = normalize_video_source(video)
+        geometry = resolve_finger_geometry(
+            export_dir=export_dir,
+            config=config,
+            map_cache=map_cache,
+            source_scale=source_scale,
+            calib=calib,
+        )
+
+        mesh_path, intrinsics_path, pose_path = geometry.export_paths.as_tuple()
+        domain = build_or_load_domain(
+            mesh_path,
+            intrinsics_path,
+            pose_path,
+            geometry.map_cache,
+            geometry.source_K,
+            geometry.source_width,
+            geometry.source_height,
+            geometry.out_width,
+            geometry.out_height,
+            edge_expand=geometry.edge_expand,
+            include_surface=True,
+            verbose=verbose,
+            source_scale=geometry.source_scale,
+        )
+        domain.remap.ensure_fast_maps()
+
+        capture = open_source(video, geometry.config)
+        if not capture.isOpened():
+            raise RuntimeError("Cannot open camera/video source")
+
+        sensor = Sensor("external", config_name=config, verbose=verbose, mirror=mirror)
+        frame_processor = UVUnrollFrameProcessor(
+            domain.remap,
+            camera_matrix=geometry.source_K if geometry.dist_coeffs is not None else None,
+            dist_coeffs=geometry.dist_coeffs,
+            source_width=geometry.source_width,
+            source_height=geometry.source_height,
+            flip_source_vertical=flip_source_vertical,
+        )
+        tactile = TactileProcessor(sensor, default_valid_mask=frame_processor.valid_mask)
+        return cls(
+            geometry=geometry,
+            domain=domain,
+            tactile=tactile,
+            frame_processor=frame_processor,
+            capture=capture,
+        )
+
+    @property
+    def valid_mask(self):
+        return self.frame_processor.valid_mask
+
+    def read_source(self):
+        return self.capture.read()
+
+    def process_frame(self, frame, *, color=False):
+        """Prepare, unroll, and feed one source frame into the tactile processor."""
+        if color:
+            prepared = self.frame_processor.prepare_source(frame)
+            unrolled = self.frame_processor.flatten_prepared(prepared)
+            self.tactile.push_frame(unrolled)
+            return unrolled
+
+        prepared_gray = self.frame_processor.prepare_source_gray(frame)
+        unrolled = self.frame_processor.flatten_prepared_gray(prepared_gray)
+        self.tactile.push_frame(unrolled, already_gray=True)
+        return unrolled
+
+    def grab_and_process(self, *, color=False):
+        ok, frame = self.read_source()
+        if not ok:
+            return False, None
+        return True, self.process_frame(frame, color=color)
+
+    def get_img(self, *, color=False):
+        """Sensor-style alias: capture, unroll, and return the tactile frame or None."""
+        ok, unrolled = self.grab_and_process(color=color)
+        return unrolled if ok else None
+
+    def warmup(self, *, color=False):
+        ok, unrolled = self.grab_and_process(color=color)
+        if ok:
+            self.tactile.warmup()
+        return ok, unrolled
+
+    def close(self):
+        if self.capture is not None:
+            self.capture.release()
+        self.tactile.disconnect()