diff --git a/dimos/mapping/utils/cli/dataset_validation.md b/dimos/mapping/utils/cli/dataset_validation.md index 5f32853ddf..79b8631b54 100644 --- a/dimos/mapping/utils/cli/dataset_validation.md +++ b/dimos/mapping/utils/cli/dataset_validation.md @@ -1,16 +1,12 @@ Dataset Validation ```sh -dimos map summary recording_go2_mid360_2026-05-29_4-45pm-PST.db - -Stream("color_image"): 11141 items, 2026-05-29 23:32:57 — 2026-05-29 23:45:57 (780.1s) -Stream("fastlio_lidar"): 7240 items, 2026-05-29 23:32:56 — 2026-05-29 23:45:57 (781.7s) -Stream("fastlio_odometry"): 18737 items, 2026-05-29 23:32:56 — 2026-05-29 23:45:57 (781.8s) -Stream("lidar"): 6025 items, 2026-05-29 23:32:55 — 2026-05-29 23:45:57 (782.3s) -Stream("odom"): 14630 items, 2026-05-29 23:32:55 — 2026-05-29 23:45:57 (782.3s) +dimos mem summary recording_go2_mid360_2026-05-29_4-45pm-PST.db ``` -Shows which streams are in the database. You can replay messages in rerun: +Prints a table of the streams in the database: per-stream item count, average +Hz, start time (UTC), duration, and stored size, plus a total row. You can +replay messages in rerun: ```sh dimos map replay recording_go2_mid360_2026-05-29_4-45pm-PST.db --duration 60 diff --git a/dimos/mapping/utils/cli/map.py b/dimos/mapping/utils/cli/map.py index 106e141797..de003bf15e 100644 --- a/dimos/mapping/utils/cli/map.py +++ b/dimos/mapping/utils/cli/map.py @@ -526,50 +526,53 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: graph: PoseGraph | None = None if pgo: print("running PGO twopass map...") - prog = progress(total, "pgo pass 1 (optimizing)") - graph = lidar.tap(prog).transform(PGO()).last().data + with progress(total, "pgo pass 1 (optimizing)") as bar: + graph = lidar.tap(bar).transform(PGO()).last().data pgo_path = [ (kf.optimized.translation.x, kf.optimized.translation.y, kf.optimized.translation.z) for kf in graph.keyframes ] - pgo_map = _accumulate( - kept, - voxel=voxel, - block_count=block_count, - device=device, - graph=graph, - register=register, - carve_columns=carve, - progress_cb=progress(n_kept, "pgo pass 2 (rebuilding)"), - ) + with progress(n_kept, "pgo pass 2 (rebuilding)") as bar: + pgo_map = _accumulate( + kept, + voxel=voxel, + block_count=block_count, + device=device, + graph=graph, + register=register, + carve_columns=carve, + progress_cb=bar, + ) full_pgo_map = None if full_pgo: assert graph is not None - full_pgo_map = _accumulate( - lidar, + with progress(total, "full pgo (rebuilding)") as bar: + full_pgo_map = _accumulate( + lidar, + voxel=voxel, + block_count=block_count, + device=device, + graph=graph, + register=register, + carve_columns=carve, + progress_cb=bar, + ) + + # Raw map: same dedup'd frames, no PGO correction. + with progress(n_kept, "reconstructing global map") as bar: + global_map = _accumulate( + kept, voxel=voxel, block_count=block_count, device=device, - graph=graph, register=register, carve_columns=carve, - progress_cb=progress(total, "full pgo (rebuilding)"), + progress_cb=bar, ) - # Raw map: same dedup'd frames, no PGO correction. - global_map = _accumulate( - kept, - voxel=voxel, - block_count=block_count, - device=device, - register=register, - carve_columns=carve, - progress_cb=progress(n_kept, "reconstructing global map"), - ) - marker_dets: list[Observation[Any]] = [] if markers: # Image observations in dimos recordings are stamped with @@ -596,17 +599,18 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: # Keep the sharpest frame per --marker-quality-window window, then # drop frames where the robot was moving (linear + rotational) faster # than the limits. Defaults match replay_marker.py so positions agree. - pipeline: Stream[Image] = color_image.tap( - progress(n_images, "detecting markers") - ).transform(QualityWindow(lambda img: img.sharpness, window=marker_quality_window)) - if marker_max_speed > 0: - pipeline = pipeline.transform( - SpeedLimit( - max_mps=marker_max_speed, - max_dps=marker_max_rot_rate if marker_max_rot_rate > 0 else None, - ) + with progress(n_images, "detecting markers") as bar: + pipeline: Stream[Image] = color_image.tap(bar).transform( + QualityWindow(lambda img: img.sharpness, window=marker_quality_window) ) - all_dets = pipeline.transform(xf).to_list() + if marker_max_speed > 0: + pipeline = pipeline.transform( + SpeedLimit( + max_mps=marker_max_speed, + max_dps=marker_max_rot_rate if marker_max_rot_rate > 0 else None, + ) + ) + all_dets = pipeline.transform(xf).to_list() if marker_smoothing > 0: # Keep only the latest emission per track_id — that's the most # averaged pose, drawn once per tracked marker session. diff --git a/dimos/mapping/utils/cli/rename.py b/dimos/mapping/utils/cli/rename.py index e856131e86..a55701ccfb 100644 --- a/dimos/mapping/utils/cli/rename.py +++ b/dimos/mapping/utils/cli/rename.py @@ -31,7 +31,7 @@ import typer -from dimos.memory2.codecs.base import _resolve_payload_type +from dimos.memory2.codecs.base import resolve_payload_type from dimos.memory2.store.sqlite import SqliteStore from dimos.memory2.stream import Stream from dimos.memory2.type.observation import Observation @@ -77,7 +77,7 @@ def _stream_payload_types(db_path: Path) -> dict[str, type]: rows = conn.execute("SELECT name, config FROM _streams").fetchall() finally: conn.close() - return {name: _resolve_payload_type(json.loads(cfg)["payload_module"]) for name, cfg in rows} + return {name: resolve_payload_type(json.loads(cfg)["payload_module"]) for name, cfg in rows} def _parse_renames(pairs: list[str]) -> dict[str, str]: diff --git a/dimos/mapping/utils/cli/replay.py b/dimos/mapping/utils/cli/replay.py index c78704f725..209ab6ee95 100644 --- a/dimos/mapping/utils/cli/replay.py +++ b/dimos/mapping/utils/cli/replay.py @@ -31,57 +31,25 @@ from __future__ import annotations -from collections.abc import Callable from pathlib import Path import subprocess -import time from typing import TYPE_CHECKING, Any import rerun as rr import typer +from dimos.memory2.utils.progress import progress + # Heavy dimos imports (mapping/memory2 → torch, scipy, open3d) are deferred into # main() so that `dimos map --help` stays fast. See test_cli_startup.py and the # same pattern in dimos/mapping/utils/cli/map.py. if TYPE_CHECKING: from dimos.memory2.stream import Stream - from dimos.memory2.type.observation import Observation from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 TIMELINE = "ts" -def _progress(total: int, label: str) -> Callable[[Observation[Any]], None]: - """Matches dimos/mapping/utils/cli/map.py:progress.""" - seen = 0 - wall_start: float | None = None - last_wall: float | None = None - first_ts: float | None = None - - def tick(obs: Observation[Any]) -> None: - nonlocal seen, wall_start, last_wall, first_ts - now = time.monotonic() - if wall_start is None: - wall_start = now - first_ts = obs.ts - assert first_ts is not None - frame_ms = (now - last_wall) * 1000 if last_wall is not None else 0.0 - last_wall = now - seen += 1 - pct = 100 * seen // total if total else 100 - wall = now - wall_start - data = obs.ts - first_ts - speed = data / wall if wall > 0 else 0.0 - end = "\n" if seen >= total else "" - print( - f"\r{label} {pct:>3}% [{seen}/{total}] {data:.1f}s ({speed:.1f} x rt) {frame_ms:.0f}ms/frame", - end=end, - flush=True, - ) - - return tick - - def _log_clouds( label: str, stream: Stream[PointCloud2], @@ -99,14 +67,14 @@ def _log_clouds( whole pipeline. """ n = total if total is not None else stream.count() - cb = _progress(n, label) - for obs in stream: - cb(obs) - rr.set_time(TIMELINE, timestamp=obs.ts) - rr.log( - entity, - obs.data.to_rerun(voxel_size=voxel, mode=point_mode, bottom_cutoff=bottom_cutoff), - ) + with progress(n, label) as bar: + for obs in stream: + bar(obs) + rr.set_time(TIMELINE, timestamp=obs.ts) + rr.log( + entity, + obs.data.to_rerun(voxel_size=voxel, mode=point_mode, bottom_cutoff=bottom_cutoff), + ) def _log_path( @@ -122,22 +90,22 @@ def _log_path( without a pose are skipped. """ n = stream.count() - cb = _progress(n, label) points: list[tuple[float, float, float]] = [] last_ts: float | None = None emit_count = 0 - for obs in stream: - cb(obs) - if obs.pose_tuple is None: - continue - points.append( - (float(obs.pose_tuple[0]), float(obs.pose_tuple[1]), float(obs.pose_tuple[2])) - ) - last_ts = obs.ts - emit_count += 1 - if emit_every > 0 and emit_count % emit_every == 0 and len(points) >= 2: - rr.set_time(TIMELINE, timestamp=obs.ts) - rr.log(entity, rr.LineStrips3D([points], colors=[color])) + with progress(n, label) as bar: + for obs in stream: + bar(obs) + if obs.pose_tuple is None: + continue + points.append( + (float(obs.pose_tuple[0]), float(obs.pose_tuple[1]), float(obs.pose_tuple[2])) + ) + last_ts = obs.ts + emit_count += 1 + if emit_every > 0 and emit_count % emit_every == 0 and len(points) >= 2: + rr.set_time(TIMELINE, timestamp=obs.ts) + rr.log(entity, rr.LineStrips3D([points], colors=[color])) if ( last_ts is not None and len(points) >= 2 @@ -206,8 +174,8 @@ def main( ), ) -> None: """Dump a recording to .rrd (lidar clouds + camera frames) and open it in rerun.""" - from dimos.mapping.utils.cli.summary import _stream_payload_types from dimos.mapping.voxels import VoxelMapTransformer + from dimos.memory2.cli.summary import stream_payload_types from dimos.memory2.store.sqlite import SqliteStore from dimos.memory2.transform import throttle from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -224,7 +192,7 @@ def main( # Resolve which streams to voxelize: all PointCloud2 streams, or the # explicit --map-source subset. Validate up front so typos fail fast. - pc_streams = [n for n, t in _stream_payload_types(db_path).items() if t is PointCloud2] + pc_streams = [n for n, t in stream_payload_types(db_path).items() if t is PointCloud2] map_sources = list(map_source) or pc_streams if (map or map_final) and (bad := [s for s in map_sources if s not in pc_streams]): raise typer.BadParameter(f"--map-source: not PointCloud2 stream(s): {', '.join(bad)}") @@ -309,20 +277,20 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: # fastlio pose axis + path from fastlio_odometry stream. if "fastlio_odometry" in store.streams: odometry = clipped("fastlio_odometry", Odometry) - cb = _progress(odometry.count(), "fastlio_odometry") - for obs in odometry: - cb(obs) - if obs.pose_tuple is None: - continue - rr.set_time(TIMELINE, timestamp=obs.ts) - x, y, z, qx, qy, qz, qw = obs.pose_tuple - rr.log( - "world/fastlio", - rr.Transform3D( - translation=[x, y, z], - quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]), - ), - ) + with progress(odometry.count(), "fastlio_odometry") as bar: + for obs in odometry: + bar(obs) + if obs.pose_tuple is None: + continue + rr.set_time(TIMELINE, timestamp=obs.ts) + x, y, z, qx, qy, qz, qw = obs.pose_tuple + rr.log( + "world/fastlio", + rr.Transform3D( + translation=[x, y, z], + quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]), + ), + ) _log_path( " fastlio_path", clipped("fastlio_odometry", Odometry), @@ -333,20 +301,20 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: # Go2 native odom pose axis + path. if "odom" in store.streams: odom = clipped("odom", PoseStamped) - cb = _progress(odom.count(), " odom") - for odom_obs in odom: - cb(odom_obs) - if odom_obs.pose_tuple is None: - continue - rr.set_time(TIMELINE, timestamp=odom_obs.ts) - x, y, z, qx, qy, qz, qw = odom_obs.pose_tuple - rr.log( - "world/odom", - rr.Transform3D( - translation=[x, y, z], - quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]), - ), - ) + with progress(odom.count(), " odom") as bar: + for odom_obs in odom: + bar(odom_obs) + if odom_obs.pose_tuple is None: + continue + rr.set_time(TIMELINE, timestamp=odom_obs.ts) + x, y, z, qx, qy, qz, qw = odom_obs.pose_tuple + rr.log( + "world/odom", + rr.Transform3D( + translation=[x, y, z], + quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]), + ), + ) _log_path( " odom_path", clipped("odom", PoseStamped), @@ -359,19 +327,19 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: color_image.transform(throttle(1.0 / camera_hz)) if camera_hz > 0 else color_image ) n_img = cam_pipeline.count() - cb = _progress(n_img, " color_image") - for img_obs in cam_pipeline: - cb(img_obs) - rr.set_time(TIMELINE, timestamp=img_obs.ts) - if img_obs.pose_tuple is not None: - x, y, z, qx, qy, qz, qw = img_obs.pose_tuple - rr.log( - "world/camera", - rr.Transform3D( - translation=[x, y, z], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]) - ), - ) - rr.log("world/camera/image", img_obs.data.to_rerun()) + with progress(n_img, " color_image") as bar: + for img_obs in cam_pipeline: + bar(img_obs) + rr.set_time(TIMELINE, timestamp=img_obs.ts) + if img_obs.pose_tuple is not None: + x, y, z, qx, qy, qz, qw = img_obs.pose_tuple + rr.log( + "world/camera", + rr.Transform3D( + translation=[x, y, z], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]) + ), + ) + rr.log("world/camera/image", img_obs.data.to_rerun()) print(f"wrote {out}") if no_gui: diff --git a/dimos/mapping/utils/cli/summary.py b/dimos/mapping/utils/cli/summary.py deleted file mode 100644 index 694565b6e8..0000000000 --- a/dimos/mapping/utils/cli/summary.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Print ``Store.summary()`` for a memory2 sqlite recording. - -Usage: - uv run dimos map summary mid360 -""" - -from __future__ import annotations - -import json -from pathlib import Path -import sqlite3 - -import typer - -from dimos.memory2.codecs.base import _resolve_payload_type -from dimos.memory2.store.sqlite import SqliteStore -from dimos.utils.data import resolve_named_path - - -def _stream_payload_types(db_path: Path) -> dict[str, type]: - """Read each stream's registered payload type from the _streams table.""" - conn = sqlite3.connect(str(db_path)) - try: - rows = conn.execute("SELECT name, config FROM _streams").fetchall() - finally: - conn.close() - return {name: _resolve_payload_type(json.loads(cfg)["payload_module"]) for name, cfg in rows} - - -def main( - dataset: str = typer.Argument(..., help="Dataset .db: bare name (cwd or data/) or path"), -) -> None: - """Print per-stream counts and time ranges for a recorded SQLite dataset.""" - db_path = resolve_named_path(dataset, ".db") - payload_types = _stream_payload_types(db_path) - - store = SqliteStore(path=str(db_path)) - with store: - for name, ptype in payload_types.items(): - store.stream(name, ptype) - print(store.summary()) - - -if __name__ == "__main__": - typer.run(main) diff --git a/dimos/mapping/utils/cli/test_cli.py b/dimos/mapping/utils/cli/test_cli.py index cb9ece1f3c..208d7b71b7 100644 --- a/dimos/mapping/utils/cli/test_cli.py +++ b/dimos/mapping/utils/cli/test_cli.py @@ -66,8 +66,8 @@ def _turbojpeg_available() -> bool: ) -def _run(*args: str, timeout: float = 300.0) -> SimpleNamespace: - """Invoke `dimos map ` in-process and capture its result. +def _run(*args: str, timeout: float = 300.0, group: str = "map") -> SimpleNamespace: + """Invoke `dimos ` in-process and capture its result. Uses Typer's CliRunner so the dimos import cost is paid once (module import) rather than per case. `timeout` is kept for call-site compatibility but is a @@ -76,7 +76,7 @@ def _run(*args: str, timeout: float = 300.0) -> SimpleNamespace: """ from dimos.robot.cli.dimos import main as cli_app - res = _runner.invoke(cli_app, ["map", *args]) + res = _runner.invoke(cli_app, [group, *args]) err = res.output if res.exception is not None and not isinstance(res.exception, SystemExit): err += "\n" + "".join(traceback.format_exception(res.exception)) @@ -103,10 +103,12 @@ def dataset() -> str: def test_summary(dataset: str) -> None: - res = _run("summary", dataset) + # summary lives under `dimos mem`, but drives the same recorded dataset. + res = _run("summary", dataset, group="mem") assert res.returncode == 0, res.stderr assert "lidar" in res.stdout assert "odom" in res.stdout + assert "iB" in res.stdout # payload sizes (KiB/MiB/GiB) are included @requires_turbojpeg diff --git a/dimos/memory2/backend.py b/dimos/memory2/backend.py index 7b95bd6335..e96075d929 100644 --- a/dimos/memory2/backend.py +++ b/dimos/memory2/backend.py @@ -76,6 +76,10 @@ def start(self) -> None: def name(self) -> str: return self.metadata_store.name + def size_bytes(self) -> int | None: + """Total stored payload bytes for this stream, or None if not cheaply knowable.""" + return self.blob_store.size_bytes(self.name) if self.blob_store is not None else None + def _make_loader(self, row_id: int) -> Any: bs = self.blob_store if bs is None: diff --git a/dimos/memory2/blobstore/base.py b/dimos/memory2/blobstore/base.py index 36601c5f10..e4fbc1a272 100644 --- a/dimos/memory2/blobstore/base.py +++ b/dimos/memory2/blobstore/base.py @@ -54,5 +54,9 @@ def delete(self, stream_name: str, key: int) -> None: """Delete a blob by stream name and observation id.""" ... + def size_bytes(self, stream_name: str) -> int | None: + """Total stored payload bytes for a stream, or None if not cheaply knowable.""" + return None + def serialize(self) -> dict[str, Any]: return {"class": qual(type(self)), "config": self.config.model_dump()} diff --git a/dimos/memory2/blobstore/file.py b/dimos/memory2/blobstore/file.py index 452b13c9b7..60c8997db0 100644 --- a/dimos/memory2/blobstore/file.py +++ b/dimos/memory2/blobstore/file.py @@ -67,3 +67,13 @@ def delete(self, stream_name: str, key: int) -> None: p.unlink() except FileNotFoundError: raise KeyError(f"No blob for stream={stream_name!r}, key={key}") from None + + def size_bytes(self, stream_name: str) -> int | None: + validate_identifier(stream_name) + stream_dir = self._root / stream_name + if not stream_dir.is_dir(): + return 0 + # Only integer-keyed files — the layout put() writes. + return sum( + p.stat().st_size for p in stream_dir.glob("*.bin") if p.is_file() and p.stem.isdigit() + ) diff --git a/dimos/memory2/blobstore/sqlite.py b/dimos/memory2/blobstore/sqlite.py index 85b539bfd6..06021dcd9a 100644 --- a/dimos/memory2/blobstore/sqlite.py +++ b/dimos/memory2/blobstore/sqlite.py @@ -105,3 +105,14 @@ def delete(self, stream_name: str, key: int) -> None: raise KeyError(f"No blob for stream={stream_name!r}, key={key}") from None if cur.rowcount == 0: raise KeyError(f"No blob for stream={stream_name!r}, key={key}") + + def size_bytes(self, stream_name: str) -> int | None: + # LENGTH(data) comes from the record header, so this never reads payload pages. + validate_identifier(stream_name) + try: + row = self._conn.execute( + f'SELECT SUM(LENGTH(data)) FROM "{stream_name}_blob"' + ).fetchone() + except sqlite3.OperationalError: # no blob table for this stream + return 0 + return int(row[0]) if row[0] is not None else 0 diff --git a/dimos/memory2/blobstore/test_blobstore.py b/dimos/memory2/blobstore/test_blobstore.py index ade6aa4cc6..2cae1a8a55 100644 --- a/dimos/memory2/blobstore/test_blobstore.py +++ b/dimos/memory2/blobstore/test_blobstore.py @@ -60,3 +60,12 @@ def test_large_blob(self, blob_store: BlobStore) -> None: blob_store.put("big", 0, data) assert blob_store.get("big", 0) == data assert blob_store.get("big", 0) == data + + def test_size_bytes(self, blob_store: BlobStore) -> None: + assert blob_store.size_bytes("s") == 0 + blob_store.put("s", 1, b"12345") + blob_store.put("s", 2, b"1234567890") + assert blob_store.size_bytes("s") == 15 + blob_store.delete("s", 1) + assert blob_store.size_bytes("s") == 10 + assert blob_store.size_bytes("other") == 0 diff --git a/dimos/memory2/cli/app.py b/dimos/memory2/cli/app.py index 2424c600df..8c4278a207 100644 --- a/dimos/memory2/cli/app.py +++ b/dimos/memory2/cli/app.py @@ -18,12 +18,15 @@ import typer +from dimos.memory2.cli.summary import main as _summary_main + mem_app = typer.Typer(help="memory2 store commands", no_args_is_help=True) +mem_app.command("summary")(_summary_main) @mem_app.command() def rerun( - path: str = typer.Argument(..., help="Store to render (.mcap path/name or .db)"), + path: str = typer.Argument(..., help="Store: bare name (cwd, data/, LFS), .db or .mcap path"), out: str = typer.Option(None, "--out", help="Output .rrd (default: alongside the source)"), seconds: float = typer.Option(None, "--seconds", help="Only the first N seconds"), no_gui: bool = typer.Option(False, "--no-gui", help="Write the .rrd but don't open the viewer"), diff --git a/dimos/memory2/cli/render.py b/dimos/memory2/cli/render.py index 0869c7c45f..fe50c0d489 100644 --- a/dimos/memory2/cli/render.py +++ b/dimos/memory2/cli/render.py @@ -28,19 +28,26 @@ import subprocess from typing import TYPE_CHECKING +from dimos.utils.data import resolve_named_path + if TYPE_CHECKING: from dimos.memory2.store.base import Store def open_store(path: str) -> Store: - """Open a store by file type (``.db`` -> SqliteStore, else Go2 mcap).""" - if str(path).endswith(".db"): - from dimos.memory2.store.sqlite import SqliteStore + """Open a store by name or path (``.db`` -> SqliteStore, ``.mcap`` -> Go2 mcap). + + Bare names resolve like the other db verbs: cwd, then ``data/``, then an + LFS pull — defaulting to ``.db`` unless the name says ``.mcap``. + """ + resolved = resolve_named_path(path, ".mcap" if str(path).endswith(".mcap") else ".db") + if resolved.suffix == ".mcap": + from dimos.robot.unitree.go2.dds.store import Go2McapStore # lazy: robot-layer codec set - return SqliteStore(path=path, must_exist=True) - from dimos.robot.unitree.go2.dds.store import Go2McapStore # lazy: robot-layer codec set + return Go2McapStore(path=str(resolved)) + from dimos.memory2.store.sqlite import SqliteStore - return Go2McapStore(path=path) + return SqliteStore(path=str(resolved), must_exist=True) def _open_viewer(rrd: str) -> None: @@ -102,22 +109,21 @@ def render_store( rr.save(out) for name, stream in renderable: - report = progress(stream.count(), label=name) - for obs in stream: - if seconds is not None and obs.ts - t0 > seconds: - print() # terminate the windowed (sub-100%) progress line - break - if obs.data is None: # e.g. a truncated/corrupt frame that failed to decode + with progress(stream.count(), label=name) as report: + for obs in stream: + if seconds is not None and obs.ts - t0 > seconds: + break # the context manager finalizes the windowed (sub-100%) bar + if obs.data is None: # e.g. a truncated/corrupt frame that failed to decode + report(obs) + continue + rr.set_time("time", duration=obs.ts - t0) + data = obs.data.to_rerun() + if isinstance(data, list): # RerunMulti: [(subpath, archetype), ...] + for sub, arch in data: + rr.log(f"{name}/{sub}", arch) + else: + rr.log(name, data) report(obs) - continue - rr.set_time("time", duration=obs.ts - t0) - data = obs.data.to_rerun() - if isinstance(data, list): # RerunMulti: [(subpath, archetype), ...] - for sub, arch in data: - rr.log(f"{name}/{sub}", arch) - else: - rr.log(name, data) - report(obs) rr.rerun_shutdown() # flush + close the .rrd before opening it print(f"wrote {out}") diff --git a/dimos/memory2/cli/summary.py b/dimos/memory2/cli/summary.py new file mode 100644 index 0000000000..8aa2519915 --- /dev/null +++ b/dimos/memory2/cli/summary.py @@ -0,0 +1,138 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Print ``Store.summary()`` for a memory2 sqlite recording. + +Usage: + uv run dimos mem summary mid360 +""" + +from __future__ import annotations + +from datetime import datetime, timezone +import json +from math import log +import sqlite3 +from typing import TYPE_CHECKING, Any + +import typer + +from dimos.utils.colors import HEAT_GRADIENT_ANSI256 +from dimos.utils.data import resolve_named_path +from dimos.utils.human import human_bytes + +if TYPE_CHECKING: + from pathlib import Path + + from dimos.memory2.stream import Stream + +# Heavy dimos imports (memory2 store → codecs, msgs) and rich are deferred into +# the function bodies so that `dimos --help` — which imports this module just to +# register the `mem summary` command — stays fast. See test_cli_startup.py. + + +def stream_payload_types(db_path: Path) -> dict[str, type]: + """Read each stream's registered payload type from the _streams table.""" + from dimos.memory2.codecs.base import resolve_payload_type + + conn = sqlite3.connect(str(db_path)) + try: + rows = conn.execute("SELECT name, config FROM _streams").fetchall() + finally: + conn.close() + return {name: resolve_payload_type(json.loads(cfg)["payload_module"]) for name, cfg in rows} + + +def _shade(value: float, lo: float, hi: float) -> str: + """Rich style for ``value`` relative to [lo, hi], log-scaled (columns span decades).""" + if value <= 0: + return "dim" + t = 0.5 if hi <= lo else (log(value) - log(lo)) / (log(hi) - log(lo)) + return f"color({HEAT_GRADIENT_ANSI256[round(t * (len(HEAT_GRADIENT_ANSI256) - 1))]})" + + +def _heat(text: str, value: float, column: list[float]) -> str: + """Wrap ``text`` in rich markup colored by ``value``'s rank within ``column``.""" + positive = [v for v in column if v > 0] + lo, hi = (min(positive), max(positive)) if positive else (0.0, 0.0) + return f"[{_shade(value, lo, hi)}]{text}[/]" + + +def main( + dataset: str = typer.Argument(..., help="Dataset .db: bare name (cwd or data/) or path"), +) -> None: + """Print per-stream counts, time ranges, and payload sizes for a recorded SQLite dataset.""" + from rich.console import Console + from rich.progress import Progress + from rich.table import Table + + from dimos.memory2.store.sqlite import SqliteStore + + db_path = resolve_named_path(dataset, ".db") + payload_types = stream_payload_types(db_path) + + rows: list[tuple[str, int, float | None, float | None, int]] = [] + store = SqliteStore(path=str(db_path)) + with store, Progress(transient=True) as prog: + task = prog.add_task("scanning", total=len(payload_types)) + for name, ptype in payload_types.items(): + prog.update(task, description=name) + stream: Stream[Any] = store.stream(name, ptype) + n = stream.count() + t0, t1 = stream.get_time_range() if n else (None, None) + rows.append((name, n, t0, t1, stream.size_bytes() or 0)) + prog.advance(task) + rows.sort(key=lambda r: r[4], reverse=True) + + table = Table(title=db_path.name) + table.add_column("Stream", style="cyan") + table.add_column("Items", justify="right") + table.add_column("Hz", justify="right") + table.add_column("Start (UTC)") + table.add_column("Duration", justify="right") + table.add_column("Size", justify="right") + + def hz(n: int, t0: float | None, t1: float | None) -> float: + return (n - 1) / (t1 - t0) if t0 is not None and t1 is not None and t1 > t0 else 0.0 + + items_col = [float(r[1]) for r in rows] + hz_col = [hz(r[1], r[2], r[3]) for r in rows] + size_col = [float(r[4]) for r in rows] + + for name, n, t0, t1, size in rows: + dur = t1 - t0 if t0 is not None and t1 is not None else None + rate = hz(n, t0, t1) + table.add_row( + name, + _heat(f"{n:,}", n, items_col), + _heat(f"{rate:.1f}", rate, hz_col) if rate > 0 else "—", + datetime.fromtimestamp(t0, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + if t0 is not None + else "—", + f"{dur:.1f}s" if dur is not None else "—", + _heat(human_bytes(size), size, size_col), + ) + table.add_section() + table.add_row( + "total", f"{sum(r[1] for r in rows):,}", "", "", "", human_bytes(sum(r[4] for r in rows)) + ) + + console = Console() + if not console.is_terminal: # piped: don't squeeze the table into the 80-col default + console = Console(width=250) + console.print(table) + + +if __name__ == "__main__": + typer.run(main) diff --git a/dimos/memory2/codecs/base.py b/dimos/memory2/codecs/base.py index 821b36b60f..0ddd384e84 100644 --- a/dimos/memory2/codecs/base.py +++ b/dimos/memory2/codecs/base.py @@ -81,7 +81,7 @@ def _class_to_id(codec: Any) -> str: return name.lower() -def _resolve_payload_type(payload_module: str) -> type[Any]: +def resolve_payload_type(payload_module: str) -> type[Any]: parts = payload_module.rsplit(".", 1) if len(parts) != 2: raise ValueError(f"Cannot resolve payload type from {payload_module!r}") @@ -104,7 +104,7 @@ def _make_one(name: str, payload_module: str, inner: Codec[Any] | None = None) - if name == "lcm": from dimos.memory2.codecs.lcm import LcmCodec - return LcmCodec(_resolve_payload_type(payload_module)) + return LcmCodec(resolve_payload_type(payload_module)) if name == "pickle": from dimos.memory2.codecs.pickle import PickleCodec diff --git a/dimos/memory2/store/sqlite.py b/dimos/memory2/store/sqlite.py index 229961a126..8afab3a714 100644 --- a/dimos/memory2/store/sqlite.py +++ b/dimos/memory2/store/sqlite.py @@ -69,11 +69,11 @@ def _open_connection(self) -> sqlite3.Connection: def _assemble_backend(self, name: str, stored: dict[str, Any]) -> Backend[Any]: """Reconstruct a Backend from a stored config dict.""" - from dimos.memory2.codecs.base import _resolve_payload_type, codec_from_id + from dimos.memory2.codecs.base import codec_from_id, resolve_payload_type payload_module = stored["payload_module"] codec = codec_from_id(stored["codec_id"], payload_module) - data_type = _resolve_payload_type(payload_module) + data_type = resolve_payload_type(payload_module) eager_blobs = stored.get("eager_blobs", False) page_size = stored.get("page_size", self.config.page_size) diff --git a/dimos/memory2/stream.py b/dimos/memory2/stream.py index c3cb1a17ae..539a002b13 100644 --- a/dimos/memory2/stream.py +++ b/dimos/memory2/stream.py @@ -26,6 +26,7 @@ from typing_extensions import TypeVar from dimos.core.resource import CompositeResource +from dimos.memory2.backend import Backend from dimos.memory2.buffer import BackpressureBuffer, KeepLast from dimos.memory2.transform import FnIterTransformer, FnTransformer, Transformer from dimos.memory2.type.filter import ( @@ -41,6 +42,7 @@ ) from dimos.memory2.type.observation import EmbeddedObservation, Observation from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.utils.human import human_bytes from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -49,7 +51,6 @@ import reactivex from reactivex.abc import DisposableBase, ObserverBase - from dimos.memory2.backend import Backend from dimos.models.embedding.base import Embedding T = TypeVar("T") @@ -496,10 +497,29 @@ def get_time_range(self) -> tuple[float, float]: last = self.last() return (first.ts, last.ts) + def size_bytes(self) -> int | None: + """Total stored payload bytes, or None if unknown or the query narrows the stream.""" + q = self._query + narrowed = ( + q.filters + or q.limit_val is not None + or q.offset_val + or q.search_vec is not None + or q.search_text is not None + ) + if narrowed: + return None + if isinstance(self._source, Backend): + return self._source.size_bytes() + return None + def summary(self) -> str: - """Return a short human-readable summary: count, time range, avg frequency.""" + """Return a short human-readable summary: count, time range, avg frequency, size.""" from datetime import datetime, timezone + sz = self.size_bytes() + size = f", {human_bytes(sz)}" if sz is not None else "" + n = self.count() if n == 0: return f"{self}: empty" @@ -510,7 +530,7 @@ def summary(self) -> str: dt1 = datetime.fromtimestamp(t1, tz=timezone.utc).strftime(fmt) dur = t1 - t0 hz = f", {(n - 1) / dur:.2f} Hz" if dur > 0 else "" - return f"{self}: {n} items, {dt0} — {dt1} ({dur:.1f}s{hz})" + return f"{self}: {n} items, {dt0} — {dt1} ({dur:.1f}s{hz}{size})" def materialize(self) -> Stream[T, O]: """Materialize into memory and return a replayable stream. diff --git a/dimos/memory2/utils/progress.py b/dimos/memory2/utils/progress.py index 29aaf8d4e4..05f1256ac1 100644 --- a/dimos/memory2/utils/progress.py +++ b/dimos/memory2/utils/progress.py @@ -12,51 +12,111 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""A per-observation progress callback, shaped for ``Stream.tap``. +"""A per-observation progress bar, shaped for ``Stream.tap``. - stream.tap(progress(stream.count(), "render")).drain() + with progress(stream.count(), "render") as bar: + stream.tap(bar).drain() -Prints a single rewritten line: percent, count, data-seconds covered, speed -relative to wall-clock (``x rt``), and per-frame latency. +Renders a rich progress bar: label, percent, count, data-seconds covered, +speed relative to wall-clock (``x rt``), and per-frame latency. On completion +the live bar is replaced by one plain persisted line. ``progress`` is a +context manager: exiting the ``with`` block finalizes the bar (and restores +the cursor), even when the stream is abandoned early (e.g. a bounded time +window) or the pipeline raises. rich allows a single live display, so keep +one bar per ``with`` block at a time. """ from __future__ import annotations -from collections.abc import Callable +from contextlib import contextmanager import time from typing import TYPE_CHECKING, Any if TYPE_CHECKING: + from collections.abc import Iterator + from dimos.memory2.type.observation import Observation +_REFRESH_S = 0.1 # refresh cap; keeps 200 Hz streams from redrawing per frame + + +class ProgressBar: + """Callable per-observation progress; ``close()`` finalizes an unfinished bar.""" + + def __init__(self, total: int, label: str = "") -> None: + # lazy: keep rich off the CLI startup path (see test_cli_startup.py) + from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + TaskProgressColumn, + TextColumn, + ) -def progress(total: int, label: str = "") -> Callable[[Observation[Any]], None]: - """Return a callback that prints streaming progress, one call per observation.""" - seen = 0 - wall_start: float | None = None - last_wall: float | None = None - first_ts: float | None = None + self._total = total + self._label = label or "processing" + # transient: close() prints its own persisted line, the live bar clears. + # auto_refresh=False: no render thread; refreshes (throttled) from the callback. + self._prog = Progress( + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TaskProgressColumn(), + MofNCompleteColumn(), + TextColumn("{task.fields[stats]}", style="dim"), + transient=True, + auto_refresh=False, + ) + self._task = self._prog.add_task(self._label, total=total, stats="") + self._closed = False + self._seen = 0 + self._wall_start: float | None = None + self._last_wall: float | None = None + self._last_refresh: float | None = None + self._first_ts: float | None = None + self._stats = "" - def _progress(obs: Observation[Any]) -> None: - nonlocal seen, wall_start, last_wall, first_ts + def __call__(self, obs: Observation[Any]) -> None: + if self._closed: + return # observations may keep flowing after close (e.g. a tap downstream) now = time.monotonic() - if wall_start is None: - wall_start = now - first_ts = obs.ts - assert first_ts is not None # narrowed by the same `if` above - frame_ms = (now - last_wall) * 1000 if last_wall is not None else 0.0 - last_wall = now - seen += 1 - pct = 100 * seen // total if total else 100 - wall = now - wall_start - data = obs.ts - first_ts + if self._wall_start is None: + self._wall_start = now + self._first_ts = obs.ts + self._prog.start() + assert self._first_ts is not None # narrowed by the same `if` above + frame_ms = (now - self._last_wall) * 1000 if self._last_wall is not None else 0.0 + self._last_wall = now + self._seen += 1 + wall = now - self._wall_start + data = obs.ts - self._first_ts speed = data / wall if wall > 0 else 0.0 - end = "\n" if seen >= total else "" - prefix = f"{label} " if label else "" - print( - f"\r{prefix}{pct:>3}% [{seen}/{total}] {data:.1f}s ({speed:.1f} x rt) {frame_ms:.0f}ms/frame", - end=end, - flush=True, - ) + self._stats = f"{data:.1f}s ({speed:.1f} x rt) {frame_ms:.0f}ms/frame" + self._prog.update(self._task, advance=1, stats=self._stats) + if self._seen >= self._total: + self.close() + elif self._last_refresh is None or now - self._last_refresh >= _REFRESH_S: + self._last_refresh = now + self._prog.refresh() + + def close(self) -> None: + """Clear the live bar and persist it as a plain line.""" + if self._closed: + return + self._closed = True + self._prog.stop() # transient: wipes the live bar + pct = 100 * self._seen // self._total if self._total else 100 + print(f"{self._label} {pct}% [{self._seen}/{self._total}] {self._stats}") + + +@contextmanager +def progress(total: int, label: str = "") -> Iterator[ProgressBar]: + """Yield a per-observation progress callback; the bar finalizes on exit. - return _progress + with progress(stream.count(), "render") as bar: + stream.tap(bar).drain() + """ + bar = ProgressBar(total, label) + try: + yield bar + finally: + bar.close() diff --git a/dimos/memory2/utils/test_progress.py b/dimos/memory2/utils/test_progress.py new file mode 100644 index 0000000000..309ca71408 --- /dev/null +++ b/dimos/memory2/utils/test_progress.py @@ -0,0 +1,75 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the ``progress`` context manager.""" + +from __future__ import annotations + +import pytest + +from dimos.memory2.type.observation import Observation +from dimos.memory2.utils.progress import progress + + +def _obs(ts: float) -> Observation[None]: + return Observation(ts=ts, _data=None) + + +def test_completed_bar_persists_one_line(capsys: pytest.CaptureFixture[str]) -> None: + with progress(3, "render") as bar: + for i in range(3): + bar(_obs(float(i))) + out = capsys.readouterr().out + assert out.count("render 100% [3/3]") == 1 + + +def test_early_exit_finalizes_partial_bar(capsys: pytest.CaptureFixture[str]) -> None: + with progress(4, "windowed") as bar: + bar(_obs(0.0)) + bar(_obs(0.1)) + out = capsys.readouterr().out + assert out.count("windowed 50% [2/4]") == 1 + + +def test_cleanup_on_exception(capsys: pytest.CaptureFixture[str]) -> None: + with pytest.raises(RuntimeError, match="boom"): + with progress(10, "crashy") as bar: + bar(_obs(0.0)) + raise RuntimeError("boom") + out = capsys.readouterr().out + assert out.count("crashy 10% [1/10]") == 1 + + +def test_calls_after_completion_are_ignored(capsys: pytest.CaptureFixture[str]) -> None: + with progress(2, "tap") as bar: + for i in range(5): # observations may keep flowing after the bar completes + bar(_obs(float(i))) + out = capsys.readouterr().out + assert out.count("tap 100% [2/2]") == 1 + + +def test_zero_total_reports_complete(capsys: pytest.CaptureFixture[str]) -> None: + with progress(0, "empty"): + pass + out = capsys.readouterr().out + assert out.count("empty 100% [0/0]") == 1 + + +def test_sequential_bars_each_persist(capsys: pytest.CaptureFixture[str]) -> None: + for label in ("first", "second"): + with progress(1, label) as bar: + bar(_obs(0.0)) + out = capsys.readouterr().out + assert out.count("first 100% [1/1]") == 1 + assert out.count("second 100% [1/1]") == 1 diff --git a/dimos/protocol/pubsub/benchmark/type.py b/dimos/protocol/pubsub/benchmark/type.py index eea43ac44c..6fb717e983 100644 --- a/dimos/protocol/pubsub/benchmark/type.py +++ b/dimos/protocol/pubsub/benchmark/type.py @@ -20,6 +20,7 @@ from typing import Any, Generic from dimos.protocol.pubsub.spec import MsgT, PubSub, TopicT +from dimos.utils.colors import HEAT_GRADIENT_ANSI256 from dimos.utils.human import human_bytes, human_duration, human_number MsgGen = Callable[[int], tuple[TopicT, MsgT]] @@ -164,29 +165,7 @@ def _print_heatmap( return min_val, max_val = min(all_vals), max(all_vals) - # ANSI 256 gradient: red -> orange -> yellow -> green - gradient = [ - 52, - 88, - 124, - 160, - 196, - 202, - 208, - 214, - 220, - 226, - 190, - 154, - 148, - 118, - 82, - 46, - 40, - 34, - ] - if not high_is_good: - gradient = gradient[::-1] + gradient = HEAT_GRADIENT_ANSI256 if high_is_good else HEAT_GRADIENT_ANSI256[::-1] def val_to_color(v: float) -> int: if v <= 0 or max_val == min_val: diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index d7615a2795..6c7cc18226 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -42,7 +42,6 @@ from dimos.mapping.utils.cli.rename import main as _map_rename_main from dimos.mapping.utils.cli.replay import main as _map_replay_main from dimos.mapping.utils.cli.replay_marker import main as _map_replay_marker_main -from dimos.mapping.utils.cli.summary import main as _map_summary_main from dimos.robot.unitree.go2.cli.go2tool import app as go2tool_app from dimos.utils.logging_config import setup_logger from dimos.visualization.rerun.constants import RerunOpenOption @@ -722,7 +721,6 @@ def dataprep_inspect( inspect(dataset, cast("Literal['lerobot', 'hdf5'] | None", output_format)) -map_app.command("summary")(_map_summary_main) map_app.command("rename")(_map_rename_main) map_app.command("pose-fill")(_map_pose_fill_main) map_app.command("replay")(_map_replay_main) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index b2f74bf327..fe3ee2a347 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -32,7 +32,7 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.resource import CompositeResource from dimos.core.stream import In, Out -from dimos.memory2.replay import Replay, resolve_db_path +from dimos.memory2.replay import Replay, ReplayStream, resolve_db_path from dimos.memory2.store.sqlite import SqliteStore from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion @@ -188,13 +188,29 @@ def set_motion_mode(self, name: str) -> None: def set_rage_mode(self, enable: bool) -> bool: return True + def _stream_name(self, *names: str) -> str: + """Return the first of ``names`` present in the dataset (stream naming + changed over time: mid360 recordings use go2_lidar/go2_odom, older ones + lidar/odom).""" + available = self.replay.list_streams() + for name in names: + if name in available: + return name + raise KeyError(f"None of {names!r} in dataset {self.dataset!r}; available: {available}") + @simple_mcache def lidar_stream(self) -> Observable[PointCloud2]: - return self.replay.streams.lidar.observable() + stream: ReplayStream[PointCloud2] = self.replay.stream( + self._stream_name("go2_lidar", "lidar") + ) + return stream.observable() @simple_mcache def odom_stream(self) -> Observable[PoseStamped]: - return self.replay.streams.odom.observable() + stream: ReplayStream[PoseStamped] = self.replay.stream( + self._stream_name("go2_odom", "odom") + ) + return stream.observable() @simple_mcache def video_stream(self) -> Observable[Image]: diff --git a/dimos/robot/unitree/go2/dds/cli/render.py b/dimos/robot/unitree/go2/dds/cli/render.py index ea494b41ef..748db33933 100644 --- a/dimos/robot/unitree/go2/dds/cli/render.py +++ b/dimos/robot/unitree/go2/dds/cli/render.py @@ -167,15 +167,16 @@ def log_path(obs: Observation[Path]) -> None: rr.log("world/leg_odom_path", obs.data.to_rerun()) src = store.streams.sportmodestate.to_time(seconds) - ( - src.tap(progress(src.count(), "leg_odom")) - .map_data(sportmode_pose) - .tap(log_pose) - .transform(throttle(0.1)) # reduce_rate: thin the path to ~10 Hz - .transform(accumulate_path) # yield the growing path each step - .tap(log_path) - .drain() - ) + with progress(src.count(), "leg_odom") as bar: + ( + src.tap(bar) + .map_data(sportmode_pose) + .tap(log_pose) + .transform(throttle(0.1)) # reduce_rate: thin the path to ~10 Hz + .transform(accumulate_path) # yield the growing path each step + .tap(log_path) + .drain() + ) def imu_odom(store: Go2McapStore, seconds: float | None) -> None: @@ -187,15 +188,16 @@ def log_path(obs: Observation[Path]) -> None: gravity = gravity_bias(store) # calibrate gravity+bias on the stationary start src = store.streams.imu.to_time(seconds) - ( - src.tap(progress(src.count(), "imu_odom")) - .scan_data((np.zeros(3), None), integrate_velocity(gravity)) # -> velocity (TwistStamped) - .scan_data((np.zeros(3), None), integrate_position) # -> position (PoseStamped) - .transform(throttle(0.1)) # thin the path after integrating at full IMU rate - .transform(accumulate_path) - .tap(log_path) - .drain() - ) + with progress(src.count(), "imu_odom") as bar: + ( + src.tap(bar) + .scan_data((np.zeros(3), None), integrate_velocity(gravity)) # -> velocity + .scan_data((np.zeros(3), None), integrate_position) # -> position (PoseStamped) + .transform(throttle(0.1)) # thin the path after integrating at full IMU rate + .transform(accumulate_path) + .tap(log_path) + .drain() + ) def lidar(store: Go2McapStore, seconds: float | None) -> None: @@ -207,7 +209,8 @@ def log_lidar(obs: Observation[PointCloud2]) -> None: src = store.streams.lidar.to_time(seconds) rr.log("world/leg_odom/lidar", LIDAR_TO_BASE.to_rerun(frameless=True), static=True) - (src.tap(progress(src.count(), "lidar")).tap(log_lidar).drain()) + with progress(src.count(), "lidar") as bar: + src.tap(bar).tap(log_lidar).drain() def _interp_pose( @@ -252,13 +255,14 @@ def log_voxels(obs: Observation[PointCloud2]) -> None: rr.log("world/world_lidar", obs.data.to_rerun()) src = store.streams.lidar.to_time(seconds) - ( - src.tap(progress(src.count(), "world_lidar")) - .map_data(to_world) # lidar cloud -> world-frame cloud - .transform(VoxelMapTransformer(emit_every=10, voxel_size=0.1)) # global voxel map - .tap(log_voxels) - .drain() - ) + with progress(src.count(), "world_lidar") as bar: + ( + src.tap(bar) + .map_data(to_world) # lidar cloud -> world-frame cloud + .transform(VoxelMapTransformer(emit_every=10, voxel_size=0.1)) # global voxel map + .tap(log_voxels) + .drain() + ) def camera(store: Go2McapStore, seconds: float | None, hz: float) -> None: @@ -275,7 +279,8 @@ def log_image(obs: Observation[Image]) -> None: rr.log("world/camera", obs.data.to_rerun()) src = store.streams.color_image.to_time(seconds) - (src.tap(progress(src.count(), "camera")).transform(throttle(1.0 / hz)).tap(log_image).drain()) + with progress(src.count(), "camera") as bar: + src.tap(bar).transform(throttle(1.0 / hz)).tap(log_image).drain() # Add a source: write a (store, seconds) -> None function and append it. diff --git a/dimos/utils/colors.py b/dimos/utils/colors.py index 294cf5d43b..8f5507ddfc 100644 --- a/dimos/utils/colors.py +++ b/dimos/utils/colors.py @@ -12,6 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +# ANSI 256 heatmap gradient: red -> orange -> yellow -> green. +HEAT_GRADIENT_ANSI256: list[int] = [ + 52, 88, 124, 160, 196, 202, 208, 214, 220, 226, 190, 154, 148, 118, 82, 46, 40, 34, +] # fmt: skip + def green(text: str) -> str: """Return the given text in green color."""