From a27411022dd5f3ce6ebb75b460376cb844699df8 Mon Sep 17 00:00:00 2001
From: Mishig <dmishig@gmail.com>
Date: Tue, 4 Mar 2025 10:53:01 +0100
Subject: [PATCH] [visualization] Ignore 2d or 3d data for now (#809)

---
 lerobot/scripts/visualize_dataset_html.py     | 13 ++-
 .../templates/visualize_dataset_template.html | 85 ++++++++++---------
 2 files changed, 58 insertions(+), 40 deletions(-)

diff --git a/lerobot/scripts/visualize_dataset_html.py b/lerobot/scripts/visualize_dataset_html.py
index a6899ce9..d5825aa6 100644
--- a/lerobot/scripts/visualize_dataset_html.py
+++ b/lerobot/scripts/visualize_dataset_html.py
@@ -158,7 +158,7 @@ def run_server(
             if major_version < 2:
                 return "Make sure to convert your LeRobotDataset to v2 & above."
 
-        episode_data_csv_str, columns = get_episode_data(dataset, episode_id)
+        episode_data_csv_str, columns, ignored_columns = get_episode_data(dataset, episode_id)
         dataset_info = {
             "repo_id": f"{dataset_namespace}/{dataset_name}",
             "num_samples": dataset.num_frames
@@ -218,6 +218,7 @@ def run_server(
             videos_info=videos_info,
             episode_data_csv_str=episode_data_csv_str,
             columns=columns,
+            ignored_columns=ignored_columns,
         )
 
     app.run(host=host, port=port)
@@ -236,6 +237,14 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
     selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] == "float32"]
     selected_columns.remove("timestamp")
 
+    ignored_columns = []
+    for column_name in selected_columns:
+        shape = dataset.features[column_name]["shape"]
+        shape_dim = len(shape)
+        if shape_dim > 1:
+            selected_columns.remove(column_name)
+            ignored_columns.append(column_name)
+
     # init header of csv with state and action names
     header = ["timestamp"]
 
@@ -291,7 +300,7 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
     csv_writer.writerows(rows)
     csv_string = csv_buffer.getvalue()
 
-    return csv_string, columns
+    return csv_string, columns, ignored_columns
 
 
 def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]:
diff --git a/lerobot/templates/visualize_dataset_template.html b/lerobot/templates/visualize_dataset_template.html
index 96ef28d3..cf9d40f1 100644
--- a/lerobot/templates/visualize_dataset_template.html
+++ b/lerobot/templates/visualize_dataset_template.html
@@ -224,49 +224,58 @@
                 </p>
             </div>
 
-            <table class="text-sm border-collapse border border-slate-700" x-show="currentFrameData">
-                <thead>
-                    <tr>
-                        <th></th>
-                        <template x-for="(_, colIndex) in Array.from({length: columns.length}, (_, index) => index)">
-                            <th class="border border-slate-700">
-                                <div class="flex gap-x-2 justify-between px-2">
-                                    <input type="checkbox" :checked="isColumnChecked(colIndex)"
-                                        @change="toggleColumn(colIndex)">
-                                    <p x-text="`${columns[colIndex].key}`"></p>
-                                </div>
-                            </th>
-                        </template>
-                    </tr>
-                </thead>
-                <tbody>
-                    <template x-for="(row, rowIndex) in rows">
-                        <tr class="odd:bg-gray-800 even:bg-gray-900">
-                            <td class="border border-slate-700">
-                                <div class="flex gap-x-2 max-w-64 font-semibold px-1 break-all">
-                                    <input type="checkbox" :checked="isRowChecked(rowIndex)"
-                                        @change="toggleRow(rowIndex)">
-                                </div>
-                            </td>
-                            <template x-for="(cell, colIndex) in row">
-                                <td x-show="cell" class="border border-slate-700">
-                                    <div class="flex gap-x-2 justify-between px-2" :class="{ 'hidden': cell.isNull }">
-                                        <div class="flex gap-x-2">
-                                            <input type="checkbox" x-model="cell.checked" @change="updateTableValues()">
-                                            <span x-text="`${!cell.isNull ? cell.label : null}`"></span>
-                                        </div>
-                                        <span class="w-14 text-right" x-text="`${!cell.isNull ? (typeof cell.value === 'number' ? cell.value.toFixed(2) : cell.value) : null}`"
-                                            :style="`color: ${cell.color}`"></span>
+            <div>
+                <table class="text-sm border-collapse border border-slate-700" x-show="currentFrameData">
+                    <thead>
+                        <tr>
+                            <th></th>
+                            <template x-for="(_, colIndex) in Array.from({length: columns.length}, (_, index) => index)">
+                                <th class="border border-slate-700">
+                                    <div class="flex gap-x-2 justify-between px-2">
+                                        <input type="checkbox" :checked="isColumnChecked(colIndex)"
+                                            @change="toggleColumn(colIndex)">
+                                        <p x-text="`${columns[colIndex].key}`"></p>
                                     </div>
-                                </td>
+                                </th>
                             </template>
                         </tr>
-                    </template>
-                </tbody>
-            </table>
+                    </thead>
+                    <tbody>
+                        <template x-for="(row, rowIndex) in rows">
+                            <tr class="odd:bg-gray-800 even:bg-gray-900">
+                                <td class="border border-slate-700">
+                                    <div class="flex gap-x-2 max-w-64 font-semibold px-1 break-all">
+                                        <input type="checkbox" :checked="isRowChecked(rowIndex)"
+                                            @change="toggleRow(rowIndex)">
+                                    </div>
+                                </td>
+                                <template x-for="(cell, colIndex) in row">
+                                    <td x-show="cell" class="border border-slate-700">
+                                        <div class="flex gap-x-2 justify-between px-2" :class="{ 'hidden': cell.isNull }">
+                                            <div class="flex gap-x-2">
+                                                <input type="checkbox" x-model="cell.checked" @change="updateTableValues()">
+                                                <span x-text="`${!cell.isNull ? cell.label : null}`"></span>
+                                            </div>
+                                            <span class="w-14 text-right" x-text="`${!cell.isNull ? (typeof cell.value === 'number' ? cell.value.toFixed(2) : cell.value) : null}`"
+                                                :style="`color: ${cell.color}`"></span>
+                                        </div>
+                                    </td>
+                                </template>
+                            </tr>
+                        </template>
+                    </tbody>
+                </table>
 
-            <div id="labels" class="hidden">
+                <div id="labels" class="hidden">
+                </div>
+
+                {% if ignored_columns|length > 0 %}
+                <div class="m-2 text-orange-700 max-w-96">
+                    Columns {{ ignored_columns }} are NOT shown since the visualizer currently does not support 2D or 3D data.
+                </div>
+                {% endif %}
             </div>
+
         </div>
     </div>