feat: activate self-hosted lidar spatial scene
This commit is contained in:
DCCONSTRUCTIONS
@@ -0,0 +1,386 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import time
|
||||
from collections import deque
|
||||
from collections.abc import Callable
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Literal
|
||||
|
||||
import numpy as np
|
||||
import rerun as rr
|
||||
from rerun import blueprint as rrb
|
||||
|
||||
from k1link.protocol.streams import (
|
||||
LegacyPointCloudFrame,
|
||||
LegacyPoseFrame,
|
||||
LioPointCloudFrame,
|
||||
LioPoseFrame,
|
||||
StreamDecodeError,
|
||||
decode_legacy_pointcloud,
|
||||
decode_legacy_pose,
|
||||
decode_lio_pcl,
|
||||
decode_lio_pose,
|
||||
)
|
||||
from k1link.viewer.foxglove_bridge import BridgeMetrics
|
||||
from k1link.viewer.messages import StreamMessage
|
||||
|
||||
PointColorMode = Literal["intensity", "height", "distance", "rgb", "class"]
|
||||
PointPalette = Literal["turbo", "viridis", "plasma", "grayscale", "custom"]
|
||||
|
||||
MAX_TRAJECTORY_POSES = 20_000
|
||||
DEFAULT_GRPC_PORT = 9876
|
||||
DEFAULT_CORS_ORIGINS = (
|
||||
"http://127.0.0.1:5173",
|
||||
"http://localhost:5173",
|
||||
"http://127.0.0.1:4173",
|
||||
"http://localhost:4173",
|
||||
"http://127.0.0.1:8000",
|
||||
"http://localhost:8000",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RerunSceneSettings:
|
||||
point_size: float = 2.5
|
||||
color_mode: PointColorMode = "intensity"
|
||||
palette: PointPalette = "turbo"
|
||||
custom_color: str = "#f7f8f4"
|
||||
accumulation_seconds: float = 12.0
|
||||
show_points: bool = True
|
||||
show_trajectory: bool = True
|
||||
show_grid: bool = True
|
||||
|
||||
def as_dict(self) -> dict[str, object]:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
SettingsProvider = Callable[[], RerunSceneSettings]
|
||||
|
||||
|
||||
class RerunBridge:
|
||||
"""Decode verified device topics into a self-hosted Rerun recording stream."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
grpc_port: int = DEFAULT_GRPC_PORT,
|
||||
metrics: BridgeMetrics | None = None,
|
||||
settings_provider: SettingsProvider | None = None,
|
||||
cors_allow_origin: tuple[str, ...] = DEFAULT_CORS_ORIGINS,
|
||||
recording_factory: Callable[[str], rr.RecordingStream] | None = None,
|
||||
) -> None:
|
||||
self.metrics = metrics or BridgeMetrics()
|
||||
self._settings_provider = settings_provider or RerunSceneSettings
|
||||
self._settings = self._settings_provider()
|
||||
self._recording = (recording_factory or rr.RecordingStream)(
|
||||
"nodedc_device_spatial"
|
||||
)
|
||||
blueprint = _blueprint(self._settings)
|
||||
self._url = self._recording.serve_grpc(
|
||||
grpc_port=grpc_port,
|
||||
default_blueprint=blueprint,
|
||||
server_memory_limit="512MiB",
|
||||
newest_first=True,
|
||||
cors_allow_origin=list(cors_allow_origin),
|
||||
)
|
||||
self._recording.send_blueprint(
|
||||
blueprint,
|
||||
make_active=True,
|
||||
make_default=True,
|
||||
)
|
||||
self._recording.log("/world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
|
||||
self._recording.log(
|
||||
"/world/sensor_pose",
|
||||
rr.TransformAxes3D(axis_length=0.45, show_frame=True),
|
||||
static=True,
|
||||
)
|
||||
self._path: deque[tuple[float, float, float]] = deque(
|
||||
maxlen=MAX_TRAJECTORY_POSES
|
||||
)
|
||||
self._last_trajectory_publish_ns = 0
|
||||
self._last_point_count = 0
|
||||
self._closed = False
|
||||
|
||||
@property
|
||||
def grpc_url(self) -> str:
|
||||
return self._url
|
||||
|
||||
def begin_session(self, metrics: BridgeMetrics | None = None) -> None:
|
||||
"""Reset session-local state while keeping the process-wide server alive."""
|
||||
if self._closed:
|
||||
raise RuntimeError("Rerun bridge is already closed")
|
||||
if metrics is not None:
|
||||
self.metrics = metrics
|
||||
self._settings = self._settings_provider()
|
||||
self._path.clear()
|
||||
self._last_trajectory_publish_ns = 0
|
||||
self._last_point_count = 0
|
||||
self._recording.set_time("stream_time", timestamp=time.time())
|
||||
self._recording.log("/world", rr.Clear(recursive=True))
|
||||
self._recording.log("/world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
|
||||
self._recording.log(
|
||||
"/world/sensor_pose",
|
||||
rr.TransformAxes3D(axis_length=0.45, show_frame=True),
|
||||
static=True,
|
||||
)
|
||||
self._recording.send_blueprint(
|
||||
_blueprint(self._settings),
|
||||
make_active=True,
|
||||
make_default=True,
|
||||
)
|
||||
|
||||
def process(self, message: StreamMessage) -> None:
|
||||
started_ns = time.monotonic_ns()
|
||||
self.metrics.received(len(message.payload))
|
||||
self._apply_latest_settings()
|
||||
self._set_message_time(message)
|
||||
|
||||
try:
|
||||
if message.topic.endswith("/lio_pcl"):
|
||||
self._publish_lio_pcl(decode_lio_pcl(message.payload))
|
||||
point_frame = True
|
||||
elif message.topic == "RealtimePointcloud":
|
||||
self._publish_legacy_pcl(decode_legacy_pointcloud(message.payload))
|
||||
point_frame = True
|
||||
elif message.topic.endswith("/lio_pose"):
|
||||
self._publish_lio_pose(decode_lio_pose(message.payload))
|
||||
point_frame = False
|
||||
elif message.topic == "RealtimePath":
|
||||
self._publish_legacy_pose(decode_legacy_pose(message.payload))
|
||||
point_frame = False
|
||||
else:
|
||||
return
|
||||
except StreamDecodeError:
|
||||
self.metrics.decode_error()
|
||||
return
|
||||
|
||||
published_ns = time.monotonic_ns()
|
||||
decode_publish_ms = (published_ns - started_ns) / 1_000_000
|
||||
if point_frame:
|
||||
self.metrics.published_pcl(
|
||||
self._last_point_count,
|
||||
published_ns,
|
||||
decode_publish_ms,
|
||||
)
|
||||
else:
|
||||
self.metrics.published_pose(
|
||||
published_ns,
|
||||
decode_publish_ms,
|
||||
len(self._path),
|
||||
)
|
||||
if message.source == "live_mqtt" and message.received_monotonic_ns is not None:
|
||||
self.metrics.record_latency(
|
||||
(published_ns - message.received_monotonic_ns) / 1_000_000
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
recording = self._recording
|
||||
try:
|
||||
recording.flush(timeout_sec=5.0)
|
||||
finally:
|
||||
recording.disconnect()
|
||||
# The Python wrapper owns the native gRPC server. Release it immediately
|
||||
# instead of waiting for the publisher thread frame to be collected.
|
||||
del self._recording
|
||||
|
||||
def _set_message_time(self, message: StreamMessage) -> None:
|
||||
self._recording.set_time("stream_time", timestamp=time.time())
|
||||
self._recording.set_time(
|
||||
"capture_time",
|
||||
timestamp=message.received_at_epoch_ns / 1_000_000_000,
|
||||
)
|
||||
self._recording.set_time("message_sequence", sequence=message.sequence)
|
||||
|
||||
def _apply_latest_settings(self) -> None:
|
||||
settings = self._settings_provider()
|
||||
if settings == self._settings:
|
||||
return
|
||||
self._settings = settings
|
||||
self._recording.send_blueprint(_blueprint(settings))
|
||||
|
||||
def _publish_lio_pcl(self, frame: LioPointCloudFrame) -> None:
|
||||
count = len(frame.points)
|
||||
positions = np.empty((count, 3), dtype=np.float32)
|
||||
intensities = np.empty(count, dtype=np.uint8)
|
||||
scaler = frame.header.scaler
|
||||
for index, point in enumerate(frame.points):
|
||||
positions[index] = point.scaled_xyz(scaler)
|
||||
intensities[index] = point.intensity
|
||||
self._publish_points(positions, intensities, rgb=None)
|
||||
|
||||
def _publish_legacy_pcl(self, frame: LegacyPointCloudFrame) -> None:
|
||||
count = len(frame.points)
|
||||
positions = np.empty((count, 3), dtype=np.float32)
|
||||
intensities = np.empty(count, dtype=np.uint8)
|
||||
rgb = np.empty((count, 3), dtype=np.uint8)
|
||||
for index, point in enumerate(frame.points):
|
||||
positions[index] = (point.x, point.y, point.z)
|
||||
intensities[index] = point.intensity
|
||||
rgb[index] = (point.r, point.g, point.b)
|
||||
self._publish_points(positions, intensities, rgb=rgb)
|
||||
|
||||
def _publish_points(
|
||||
self,
|
||||
positions: np.ndarray,
|
||||
intensities: np.ndarray,
|
||||
*,
|
||||
rgb: np.ndarray | None,
|
||||
) -> None:
|
||||
self._last_point_count = int(positions.shape[0])
|
||||
if not self._settings.show_points:
|
||||
self._recording.log("/world/points", rr.Clear(recursive=False))
|
||||
return
|
||||
colors = _point_colors(positions, intensities, rgb, self._settings)
|
||||
self._recording.log(
|
||||
"/world/points",
|
||||
rr.Points3D(
|
||||
positions,
|
||||
colors=colors,
|
||||
radii=rr.Radius.ui_points(self._settings.point_size),
|
||||
),
|
||||
)
|
||||
|
||||
def _publish_lio_pose(self, frame: LioPoseFrame) -> None:
|
||||
self._publish_pose(frame.position_xyz, frame.orientation_xyzw)
|
||||
|
||||
def _publish_legacy_pose(self, frame: LegacyPoseFrame) -> None:
|
||||
self._publish_pose(frame.position_xyz, frame.orientation_xyzw)
|
||||
|
||||
def _publish_pose(
|
||||
self,
|
||||
position_xyz: tuple[float, float, float],
|
||||
orientation_xyzw: tuple[float, float, float, float],
|
||||
) -> None:
|
||||
self._recording.log(
|
||||
"/world/sensor_pose",
|
||||
rr.Transform3D(
|
||||
translation=position_xyz,
|
||||
quaternion=rr.Quaternion(xyzw=orientation_xyzw),
|
||||
),
|
||||
)
|
||||
self._path.append(position_xyz)
|
||||
if not self._settings.show_trajectory:
|
||||
self._recording.log("/world/trajectory", rr.Clear(recursive=False))
|
||||
return
|
||||
now_ns = time.monotonic_ns()
|
||||
if (
|
||||
len(self._path) > 2
|
||||
and len(self._path) % 20 != 0
|
||||
and now_ns - self._last_trajectory_publish_ns < 200_000_000
|
||||
):
|
||||
return
|
||||
self._last_trajectory_publish_ns = now_ns
|
||||
self._recording.log(
|
||||
"/world/trajectory",
|
||||
rr.LineStrips3D(
|
||||
[list(self._path)],
|
||||
colors=[247, 248, 244, 255],
|
||||
radii=rr.Radius.ui_points(2.0),
|
||||
),
|
||||
)
|
||||
|
||||
def _blueprint(settings: RerunSceneSettings) -> rrb.Blueprint:
|
||||
accumulation = max(0.0, settings.accumulation_seconds)
|
||||
time_range = rr.VisibleTimeRange(
|
||||
"stream_time",
|
||||
start=rr.TimeRangeBoundary.cursor_relative(seconds=-accumulation),
|
||||
end=rr.TimeRangeBoundary.cursor_relative(),
|
||||
)
|
||||
return rrb.Blueprint(
|
||||
rrb.Spatial3DView(
|
||||
origin="/world",
|
||||
name="Пространственная сцена",
|
||||
background=[7, 8, 10, 255],
|
||||
line_grid=rrb.LineGrid3D(
|
||||
visible=settings.show_grid,
|
||||
color=[86, 91, 99, 110],
|
||||
stroke_width=0.75,
|
||||
),
|
||||
time_ranges=[time_range],
|
||||
),
|
||||
auto_layout=False,
|
||||
auto_views=False,
|
||||
collapse_panels=True,
|
||||
)
|
||||
|
||||
|
||||
def _point_colors(
|
||||
positions: np.ndarray,
|
||||
intensities: np.ndarray,
|
||||
rgb: np.ndarray | None,
|
||||
settings: RerunSceneSettings,
|
||||
) -> np.ndarray:
|
||||
if settings.color_mode == "rgb" and rgb is not None:
|
||||
return rgb
|
||||
if settings.color_mode == "class":
|
||||
color = _parse_hex_color(settings.custom_color)
|
||||
return np.tile(np.asarray(color, dtype=np.uint8), (positions.shape[0], 1))
|
||||
if settings.color_mode == "height":
|
||||
values = _normalize(positions[:, 2])
|
||||
elif settings.color_mode == "distance":
|
||||
values = _normalize(np.linalg.norm(positions, axis=1))
|
||||
else:
|
||||
values = intensities.astype(np.float32) / 255.0
|
||||
return _apply_palette(values, settings.palette, settings.custom_color)
|
||||
|
||||
|
||||
def _normalize(values: np.ndarray) -> np.ndarray:
|
||||
if values.size == 0:
|
||||
return values.astype(np.float32)
|
||||
minimum = float(np.min(values))
|
||||
maximum = float(np.max(values))
|
||||
if not math.isfinite(minimum) or not math.isfinite(maximum) or maximum <= minimum:
|
||||
return np.full(values.shape, 0.5, dtype=np.float32)
|
||||
return ((values - minimum) / (maximum - minimum)).astype(np.float32)
|
||||
|
||||
|
||||
def _apply_palette(values: np.ndarray, palette: PointPalette, custom: str) -> np.ndarray:
|
||||
clipped = np.clip(values, 0.0, 1.0)
|
||||
if palette == "custom":
|
||||
color = _parse_hex_color(custom)
|
||||
return np.tile(np.asarray(color, dtype=np.uint8), (values.size, 1))
|
||||
if palette == "grayscale":
|
||||
channel = np.rint(clipped * 255).astype(np.uint8)
|
||||
return np.column_stack((channel, channel, channel))
|
||||
|
||||
stops = {
|
||||
"turbo": (
|
||||
(0.00, (48, 18, 59)),
|
||||
(0.25, (40, 126, 231)),
|
||||
(0.50, (42, 240, 154)),
|
||||
(0.75, (246, 214, 55)),
|
||||
(1.00, (180, 4, 38)),
|
||||
),
|
||||
"viridis": (
|
||||
(0.00, (68, 1, 84)),
|
||||
(0.25, (59, 82, 139)),
|
||||
(0.50, (33, 145, 140)),
|
||||
(0.75, (94, 201, 98)),
|
||||
(1.00, (253, 231, 37)),
|
||||
),
|
||||
"plasma": (
|
||||
(0.00, (13, 8, 135)),
|
||||
(0.25, (126, 3, 168)),
|
||||
(0.50, (204, 71, 120)),
|
||||
(0.75, (248, 149, 64)),
|
||||
(1.00, (240, 249, 33)),
|
||||
),
|
||||
}[palette]
|
||||
positions = np.asarray([item[0] for item in stops], dtype=np.float32)
|
||||
colors = np.asarray([item[1] for item in stops], dtype=np.float32)
|
||||
channels = [np.interp(clipped, positions, colors[:, index]) for index in range(3)]
|
||||
return np.rint(np.column_stack(channels)).astype(np.uint8)
|
||||
|
||||
|
||||
def _parse_hex_color(value: str) -> tuple[int, int, int]:
|
||||
candidate = value.removeprefix("#")
|
||||
if len(candidate) != 6:
|
||||
return 247, 248, 244
|
||||
try:
|
||||
return tuple(int(candidate[index : index + 2], 16) for index in (0, 2, 4)) # type: ignore[return-value]
|
||||
except ValueError:
|
||||
return 247, 248, 244
|
||||
Reference in New Issue
Block a user