diff --git a/lerobot/scripts/visualize_dataset_html.py b/lerobot/scripts/visualize_dataset_html.py index 2531fbd0..7048e7a9 100644 --- a/lerobot/scripts/visualize_dataset_html.py +++ b/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, diff --git a/lerobot/templates/visualize_dataset_template.html b/lerobot/templates/visualize_dataset_template.html index e28473c1..0e762568 100644 --- a/lerobot/templates/visualize_dataset_template.html +++ b/lerobot/templates/visualize_dataset_template.html @@ -100,6 +100,13 @@ {% endfor %} + + {% if videos_info[0].language_instruction %} +

+ Language Instruction: {{ videos_info[0].language_instruction }} +

+ {% endif %} +