Add the possibility to visualize language instructions in visualize_dataset_html.py (#388)

Co-authored-by: Mishig <dmishig@gmail.com>

lerobot/scripts/visualize_dataset_html.py

@@ -112,10 +112,14 @@ def run_server(
             "fps": dataset.fps,
         }
         video_paths = get_episode_video_paths(dataset, episode_id)
+        language_instruction = get_episode_language_instruction(dataset, episode_id)
         videos_info = [
             {"url": url_for("static", filename=video_path), "filename": Path(video_path).name}
             for video_path in video_paths
         ]
+        if language_instruction:
+            videos_info[0]["language_instruction"] = language_instruction
+
         ep_csv_url = url_for("static", filename=get_ep_csv_fname(episode_id))
         return render_template(
             "visualize_dataset_template.html",
@@ -186,6 +190,20 @@ def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]
     ]
 
 
+def get_episode_language_instruction(dataset: LeRobotDataset, ep_index: int) -> list[str]:
+    # check if the dataset has language instructions
+    if "language_instruction" not in dataset.hf_dataset.features:
+        return None
+
+    # get first frame index
+    first_frame_idx = dataset.episode_data_index["from"][ep_index].item()
+
+    language_instruction = dataset.hf_dataset[first_frame_idx]["language_instruction"]
+    # TODO (michel-aractingi) hack to get the sentence, some strings in openx are badly stored
+    # with the tf.tensor appearing in the string
+    return language_instruction.removeprefix("tf.Tensor(b'").removesuffix("', shape=(), dtype=string)")
+
+
 def visualize_dataset_html(
     repo_id: str,
     root: Path | None = None,

lerobot/templates/visualize_dataset_template.html

@@ -100,6 +100,13 @@
             {% endfor %}
         </div>
 
+        <!-- Language instruction -->
+        {% if videos_info[0].language_instruction %}
+        <p class="font-medium mt-2">
+            Language Instruction: <span class="italic">{{ videos_info[0].language_instruction }}</span>
+        </p>
+        {% endif %}
+
         <!-- Shortcuts info -->
         <div class="text-sm hidden md:block">
             Hotkeys: <span class="font-mono">Space</span> to pause/unpause, <span class="font-mono">Arrow Down</span> to go to next episode, <span class="font-mono">Arrow Up</span> to go to previous episode.