From 1dafb4acf6a6c451a9f220f0432d02ef6c654f2d Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Wed, 20 May 2026 16:22:34 +0200 Subject: [PATCH] feat(viz): render depth observations as rr.DepthImage in Viridis --- src/lerobot/utils/visualization_utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index d9d5bf6b5..f13d55e0f 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -107,7 +107,14 @@ def log_rerun_data( for i, vi in enumerate(arr): rr.log(f"{key}_{i}", rr.Scalars(float(vi))) else: - img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) + if arr.shape[-1] == 1: + img_entity = ( + rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis).compress() + if compress_images + else rr.rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis) + ) + else: + img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) rr.log(key, entity=img_entity, static=True) if action: