feat(viewer): add recorded point colors and trajectory follow

src/k1link/device_plugins/xgrids_k1/observation.py

@@ -11,6 +11,9 @@ from k1link.device_plugins.xgrids_k1.archive import (
     LegacySessionCandidate,
     discover_legacy_viewer_sessions,
 )
+from k1link.device_plugins.xgrids_k1.recorded_point_colors import (
+    RecordedPointColorOverlayStore,
+)
 from k1link.device_plugins.xgrids_k1.rrd_export import (
     RrdExportCancelled,
     RrdExportError,
@@ -45,6 +48,7 @@ def build_xgrids_k1_observation(repository_root: Path) -> ObservationRuntimeCont
             resolve_missioncore_evidence_dir(repository_root),
         ),
     )
+    point_colors = RecordedPointColorOverlayStore()
     return ObservationRuntimeContribution(
         archives=tuple(
             ObservationArchiveSource(
@@ -57,6 +61,7 @@ def build_xgrids_k1_observation(repository_root: Path) -> ObservationRuntimeCont
             for archive_id, root in roots
         ),
         recording_exporter=_export_recording,
+        point_color_renderer=point_colors.render,
     )
 
 

src/k1link/device_plugins/xgrids_k1/recorded_point_colors.py

@@ -0,0 +1,434 @@
+"""On-demand point-color overlays for prepared K1 recordings."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import threading
+from collections import OrderedDict
+from contextlib import suppress
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import numpy as np
+import rerun as rr
+
+from k1link.data_plane import DecodedPointCloudView, NormalizationError
+from k1link.device_plugins.xgrids_k1.mqtt.capture import (
+    CaptureFormatError,
+    read_capture_clock_envelope,
+    read_capture_clock_origin,
+)
+from k1link.device_plugins.xgrids_k1.protocol.normalizer import normalize_k1_message
+from k1link.device_plugins.xgrids_k1.rrd_export import (
+    APPLICATION_ID,
+    CAPTURE_TIMELINE,
+    SESSION_TIMELINE,
+    _recorded_view_points,
+    _should_publish_recorded_point_frame,
+)
+from k1link.device_plugins.xgrids_k1.viewer.messages import StreamMessage
+from k1link.sessions import ReplayCommand
+from k1link.viewer.rerun_bridge import RerunSceneSettings, _point_colors
+
+_POINT_TOPICS = frozenset({"RealtimePointcloud", "lixel/application/report/lio_pcl"})
+_MAX_METADATA_LINE_BYTES = 1024 * 1024
+_MAX_PAYLOAD_BYTES = 256 * 1024 * 1024
+_DEFAULT_INDEX_CACHE_BYTES = 384 * 1024 * 1024
+_DEFAULT_PAYLOAD_CACHE_BYTES = 192 * 1024 * 1024
+
+
+class RecordedPointColorError(RuntimeError):
+    """A recorded point-color overlay could not be prepared safely."""
+
+
+@dataclass(frozen=True, slots=True)
+class _PointColorFrame:
+    sequence: int
+    session_time_ns: int
+    capture_time_ns: int
+    positions: np.ndarray
+    intensities: np.ndarray
+    rgb: np.ndarray | None
+
+    @property
+    def byte_length(self) -> int:
+        return (
+            int(self.positions.nbytes)
+            + int(self.intensities.nbytes)
+            + (0 if self.rgb is None else int(self.rgb.nbytes))
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class _PointColorIndex:
+    source_identity: tuple[object, ...]
+    frames: tuple[_PointColorFrame, ...]
+    byte_length: int
+
+
+class RecordedPointColorOverlayStore:
+    """Build small color-only RRD streams and retain bounded reusable inputs.
+
+    The base operator RRD already owns point positions. These overlays log only
+    ``Points3D:colors`` at the exact same session timestamps, so changing a
+    palette never duplicates the point geometry, trajectory, video, or AI data.
+    """
+
+    def __init__(
+        self,
+        *,
+        index_cache_bytes: int = _DEFAULT_INDEX_CACHE_BYTES,
+        payload_cache_bytes: int = _DEFAULT_PAYLOAD_CACHE_BYTES,
+    ) -> None:
+        if index_cache_bytes < 0 or payload_cache_bytes < 0:
+            raise ValueError("recorded color cache budgets must be non-negative")
+        self._index_cache_bytes = index_cache_bytes
+        self._payload_cache_bytes = payload_cache_bytes
+        self._lock = threading.RLock()
+        self._session_locks: dict[str, threading.Lock] = {}
+        self._indexes: OrderedDict[tuple[object, ...], _PointColorIndex] = OrderedDict()
+        self._index_bytes = 0
+        self._payloads: OrderedDict[tuple[object, ...], bytes] = OrderedDict()
+        self._payload_bytes = 0
+
+    def render(
+        self,
+        command: ReplayCommand,
+        *,
+        application_id: str,
+        recording_id: str,
+        color_mode: Literal["intensity", "height", "distance", "rgb", "class"],
+        palette: Literal["turbo", "viridis", "plasma", "grayscale", "custom"],
+        custom_color: str,
+    ) -> bytes:
+        if command.plugin_id != "nodedc.device.xgrids-lixelkity-k1":
+            raise RecordedPointColorError("recorded color provider received another plugin")
+        if application_id != APPLICATION_ID:
+            raise ValueError("recorded color application id is invalid")
+        if not recording_id or len(recording_id) > 128:
+            raise ValueError("recorded color recording id is invalid")
+
+        source = command.primary_artifact.path.expanduser().resolve()
+        metadata = _artifact_path(command, "raw-transport-index")
+        capture_clock = _artifact_path(command, "raw-transport-clock")
+        capture_clock_origin = _artifact_path(command, "raw-transport-clock-origin")
+        source_identity = _source_identity(
+            source,
+            metadata,
+            capture_clock,
+            capture_clock_origin,
+        )
+        settings = RerunSceneSettings(
+            color_mode=color_mode,
+            palette=palette,
+            custom_color=custom_color,
+        )
+        settings_key = (
+            color_mode,
+            palette,
+            (
+                custom_color.casefold()
+                if palette == "custom" or color_mode == "class"
+                else "-"
+            ),
+        )
+        payload_key = (*source_identity, application_id, recording_id, *settings_key)
+
+        with self._lock:
+            cached_payload = self._payloads.get(payload_key)
+            if cached_payload is not None:
+                self._payloads.move_to_end(payload_key)
+                return cached_payload
+            session_lock = self._session_locks.setdefault(command.session_id, threading.Lock())
+
+        with session_lock:
+            with self._lock:
+                cached_payload = self._payloads.get(payload_key)
+                if cached_payload is not None:
+                    self._payloads.move_to_end(payload_key)
+                    return cached_payload
+            index = self._index(
+                source_identity,
+                source,
+                metadata,
+                capture_clock,
+                capture_clock_origin,
+            )
+            payload = _render_color_overlay(
+                index.frames,
+                application_id=application_id,
+                recording_id=recording_id,
+                settings=settings,
+            )
+            self._remember_payload(payload_key, payload)
+            return payload
+
+    def _index(
+        self,
+        source_identity: tuple[object, ...],
+        source: Path,
+        metadata: Path | None,
+        capture_clock: Path | None,
+        capture_clock_origin: Path | None,
+    ) -> _PointColorIndex:
+        with self._lock:
+            cached = self._indexes.get(source_identity)
+            if cached is not None:
+                self._indexes.move_to_end(source_identity)
+                return cached
+        if metadata is None:
+            raise RecordedPointColorError("native point-color index is unavailable")
+        index = _build_index(
+            source,
+            metadata,
+            capture_clock,
+            capture_clock_origin,
+            source_identity=source_identity,
+        )
+        if index.byte_length > self._index_cache_bytes:
+            return index
+        with self._lock:
+            existing = self._indexes.pop(source_identity, None)
+            if existing is not None:
+                self._index_bytes -= existing.byte_length
+            self._indexes[source_identity] = index
+            self._index_bytes += index.byte_length
+            while self._indexes and self._index_bytes > self._index_cache_bytes:
+                _, stale = self._indexes.popitem(last=False)
+                self._index_bytes -= stale.byte_length
+        return index
+
+    def _remember_payload(self, key: tuple[object, ...], payload: bytes) -> None:
+        if len(payload) > self._payload_cache_bytes:
+            return
+        with self._lock:
+            existing = self._payloads.pop(key, None)
+            if existing is not None:
+                self._payload_bytes -= len(existing)
+            self._payloads[key] = payload
+            self._payload_bytes += len(payload)
+            while self._payloads and self._payload_bytes > self._payload_cache_bytes:
+                _, stale = self._payloads.popitem(last=False)
+                self._payload_bytes -= len(stale)
+
+
+def _build_index(
+    source: Path,
+    metadata: Path,
+    capture_clock: Path | None,
+    capture_clock_origin: Path | None,
+    *,
+    source_identity: tuple[object, ...],
+) -> _PointColorIndex:
+    try:
+        envelope = None if capture_clock is None else read_capture_clock_envelope(capture_clock)
+        origin = (
+            None
+            if capture_clock_origin is None
+            else read_capture_clock_origin(capture_clock_origin)
+        )
+    except CaptureFormatError as exc:
+        raise RecordedPointColorError("native point-color clock is invalid") from exc
+    if envelope is not None and origin is not None and (
+        envelope.started_at_epoch_ns != origin.started_at_epoch_ns
+        or envelope.started_monotonic_ns != origin.started_monotonic_ns
+    ):
+        raise RecordedPointColorError("native point-color clocks do not match")
+    session_origin_ns = (
+        envelope.started_monotonic_ns
+        if envelope is not None
+        else origin.started_monotonic_ns
+        if origin is not None
+        else None
+    )
+
+    frames: list[_PointColorFrame] = []
+    total_bytes = 0
+    point_frame_number = 0
+    previous_sequence = 0
+    previous_monotonic_ns: int | None = None
+    try:
+        source_size = source.stat().st_size
+        with source.open("rb") as raw_stream, metadata.open("r", encoding="utf-8") as index_stream:
+            for raw_line in index_stream:
+                if len(raw_line.encode("utf-8")) > _MAX_METADATA_LINE_BYTES:
+                    raise RecordedPointColorError("native point-color metadata line is too large")
+                record = json.loads(raw_line)
+                sequence = record.get("sequence")
+                monotonic_ns = record.get("received_monotonic_ns")
+                capture_time_ns = record.get("received_at_epoch_ns")
+                if (
+                    record.get("record_type") != "message"
+                    or not isinstance(sequence, int)
+                    or sequence <= previous_sequence
+                    or not isinstance(monotonic_ns, int)
+                    or monotonic_ns < 0
+                    or not isinstance(capture_time_ns, int)
+                    or capture_time_ns < 0
+                ):
+                    raise RecordedPointColorError("native point-color metadata is invalid")
+                if previous_monotonic_ns is not None and monotonic_ns < previous_monotonic_ns:
+                    raise RecordedPointColorError("native point-color timeline decreases")
+                previous_sequence = sequence
+                previous_monotonic_ns = monotonic_ns
+                if session_origin_ns is None:
+                    session_origin_ns = monotonic_ns
+                topic = record.get("topic")
+                if topic not in _POINT_TOPICS:
+                    continue
+                point_frame_number += 1
+                if not _should_publish_recorded_point_frame(point_frame_number):
+                    continue
+                payload_offset = record.get("raw_payload_offset")
+                payload_bytes = record.get("payload_bytes")
+                payload_sha256 = record.get("payload_sha256")
+                if (
+                    not isinstance(payload_offset, int)
+                    or payload_offset < 8
+                    or not isinstance(payload_bytes, int)
+                    or payload_bytes < 0
+                    or payload_bytes > _MAX_PAYLOAD_BYTES
+                    or payload_offset + payload_bytes > source_size
+                    or not isinstance(payload_sha256, str)
+                    or len(payload_sha256) != 64
+                ):
+                    raise RecordedPointColorError("native point-color offset is invalid")
+                raw_stream.seek(payload_offset)
+                payload = raw_stream.read(payload_bytes)
+                if (
+                    len(payload) != payload_bytes
+                    or hashlib.sha256(payload).hexdigest() != payload_sha256
+                ):
+                    raise RecordedPointColorError("native point-color payload is corrupt")
+                message = StreamMessage(
+                    sequence=sequence,
+                    topic=topic,
+                    payload=payload,
+                    received_at_epoch_ns=capture_time_ns,
+                    received_monotonic_ns=monotonic_ns,
+                    source="k1mqtt",
+                )
+                try:
+                    decoded = normalize_k1_message(
+                        message,
+                        processing_started_monotonic_ns=monotonic_ns,
+                    )
+                except NormalizationError as exc:
+                    raise RecordedPointColorError(
+                        "native point-color frame failed normalization"
+                    ) from exc
+                if not isinstance(decoded, DecodedPointCloudView):
+                    raise RecordedPointColorError("native point-color topic is not a point frame")
+                positions = np.asarray(decoded.positions_xyz, dtype=np.float32).reshape((-1, 3))
+                intensities = (
+                    np.full(decoded.point_count, 255, dtype=np.uint8)
+                    if decoded.intensities is None
+                    else np.frombuffer(decoded.intensities, dtype=np.uint8)
+                )
+                rgb = (
+                    None
+                    if decoded.colors_rgb is None
+                    else np.frombuffer(decoded.colors_rgb, dtype=np.uint8).reshape((-1, 3))
+                )
+                positions, intensities, rgb = _recorded_view_points(
+                    positions,
+                    intensities,
+                    rgb,
+                )
+                frame = _PointColorFrame(
+                    sequence=sequence,
+                    session_time_ns=monotonic_ns - session_origin_ns,
+                    capture_time_ns=capture_time_ns,
+                    positions=positions.copy(),
+                    intensities=intensities.copy(),
+                    rgb=None if rgb is None else rgb.copy(),
+                )
+                frames.append(frame)
+                total_bytes += frame.byte_length
+    except (OSError, json.JSONDecodeError, UnicodeError) as exc:
+        raise RecordedPointColorError("native point-color index could not be read") from exc
+    if session_origin_ns is None or not frames:
+        raise RecordedPointColorError("native point-color index contains no point frames")
+    return _PointColorIndex(
+        source_identity=source_identity,
+        frames=tuple(frames),
+        byte_length=total_bytes,
+    )
+
+
+def _render_color_overlay(
+    frames: tuple[_PointColorFrame, ...],
+    *,
+    application_id: str,
+    recording_id: str,
+    settings: RerunSceneSettings,
+) -> bytes:
+    recording = rr.RecordingStream(
+        application_id,
+        recording_id=recording_id,
+        send_properties=False,
+    )
+    stream = rr.binary_stream(recording)
+    try:
+        for frame in frames:
+            recording.set_time(
+                SESSION_TIMELINE,
+                duration=np.timedelta64(frame.session_time_ns, "ns"),
+            )
+            recording.set_time(
+                CAPTURE_TIMELINE,
+                timestamp=np.datetime64(frame.capture_time_ns, "ns"),
+            )
+            recording.set_time("message_sequence", sequence=frame.sequence)
+            recording.log(
+                "/world/points",
+                rr.Points3D.from_fields(
+                    colors=_point_colors(
+                        frame.positions,
+                        frame.intensities,
+                        frame.rgb,
+                        settings,
+                    )
+                ),
+            )
+        payload = stream.read(flush=True, flush_timeout_sec=30.0)
+    except Exception as exc:
+        raise RecordedPointColorError("recorded point-color RRD serialization failed") from exc
+    finally:
+        with suppress(Exception):
+            recording.disconnect()
+    if not payload or not payload.startswith(b"RRF2"):
+        raise RecordedPointColorError("recorded point-color RRD is invalid")
+    return payload
+
+
+def _artifact_path(command: ReplayCommand, artifact_id: str) -> Path | None:
+    matches = tuple(
+        artifact.path.expanduser().resolve()
+        for artifact in command.artifacts
+        if artifact.artifact_id == artifact_id
+    )
+    if not matches:
+        return None
+    if len(matches) != 1:
+        raise RecordedPointColorError("recorded point-color artifact is ambiguous")
+    return matches[0]
+
+
+def _source_identity(source: Path, *companions: Path | None) -> tuple[object, ...]:
+    values: list[object] = []
+    for path in (source, *companions):
+        if path is None:
+            values.extend((None, None, None))
+            continue
+        try:
+            metadata = path.stat()
+        except OSError as exc:
+            raise RecordedPointColorError("recorded point-color source is unavailable") from exc
+        if not os.path.isfile(path):
+            raise RecordedPointColorError("recorded point-color source is not a file")
+        values.extend((str(path), metadata.st_size, metadata.st_mtime_ns))
+    return tuple(values)