From 8cc7ec87c938254b297b043f03d60443fa8876e7 Mon Sep 17 00:00:00 2001 From: Mishig Davaadorj Date: Tue, 10 Jun 2025 11:39:53 +0200 Subject: [PATCH] wip --- lerobot/common/datasets/utils.py | 55 -- .../[org]/[dataset]/[episode]/fetch-data.ts | 16 +- .../src/app/[org]/[dataset]/page.tsx | 8 +- .../html_dataset_visualizer/src/app/page.tsx | 9 + lerobot/scripts/visualize_dataset_html.py | 480 ++++++------------ 5 files changed, 174 insertions(+), 394 deletions(-) diff --git a/lerobot/common/datasets/utils.py b/lerobot/common/datasets/utils.py index 542daf2aa..e8d85f0ce 100644 --- a/lerobot/common/datasets/utils.py +++ b/lerobot/common/datasets/utils.py @@ -696,61 +696,6 @@ def create_lerobot_dataset_card( ) -class IterableNamespace(SimpleNamespace): - """ - A namespace object that supports both dictionary-like iteration and dot notation access. - Automatically converts nested dictionaries into IterableNamespaces. - - This class extends SimpleNamespace to provide: - - Dictionary-style iteration over keys - - Access to items via both dot notation (obj.key) and brackets (obj["key"]) - - Dictionary-like methods: items(), keys(), values() - - Recursive conversion of nested dictionaries - - Args: - dictionary: Optional dictionary to initialize the namespace - **kwargs: Additional keyword arguments passed to SimpleNamespace - - Examples: - >>> data = {"name": "Alice", "details": {"age": 25}} - >>> ns = IterableNamespace(data) - >>> ns.name - 'Alice' - >>> ns.details.age - 25 - >>> list(ns.keys()) - ['name', 'details'] - >>> for key, value in ns.items(): - ... print(f"{key}: {value}") - name: Alice - details: IterableNamespace(age=25) - """ - - def __init__(self, dictionary: dict[str, Any] = None, **kwargs): - super().__init__(**kwargs) - if dictionary is not None: - for key, value in dictionary.items(): - if isinstance(value, dict): - setattr(self, key, IterableNamespace(value)) - else: - setattr(self, key, value) - - def __iter__(self) -> Iterator[str]: - return iter(vars(self)) - - def __getitem__(self, key: str) -> Any: - return vars(self)[key] - - def items(self): - return vars(self).items() - - def values(self): - return vars(self).values() - - def keys(self): - return vars(self).keys() - - def validate_frame(frame: dict, features: dict): expected_features = set(features) - set(DEFAULT_FEATURES) actual_features = set(frame) diff --git a/lerobot/html_dataset_visualizer/src/app/[org]/[dataset]/[episode]/fetch-data.ts b/lerobot/html_dataset_visualizer/src/app/[org]/[dataset]/[episode]/fetch-data.ts index 35c8c96f2..b25a51ee1 100644 --- a/lerobot/html_dataset_visualizer/src/app/[org]/[dataset]/[episode]/fetch-data.ts +++ b/lerobot/html_dataset_visualizer/src/app/[org]/[dataset]/[episode]/fetch-data.ts @@ -33,11 +33,17 @@ export async function getEpisodeData( }; // Generate list of episodes - const episodes = Array.from( - { length: datasetInfo.total_episodes }, - // episode id starts from 0 - (_, i) => i, - ); + const episodes = + process.env.EPISODES === undefined + ? Array.from( + { length: datasetInfo.total_episodes }, + // episode id starts from 0 + (_, i) => i, + ) + : process.env.EPISODES + .split(/\s+/) + .map((x) => parseInt(x.trim(), 10)) + .filter((x) => !isNaN(x)); // Videos information const videosInfo = Object.entries(info.features) diff --git a/lerobot/html_dataset_visualizer/src/app/[org]/[dataset]/page.tsx b/lerobot/html_dataset_visualizer/src/app/[org]/[dataset]/page.tsx index 11ed2017e..6f3109dc8 100644 --- a/lerobot/html_dataset_visualizer/src/app/[org]/[dataset]/page.tsx +++ b/lerobot/html_dataset_visualizer/src/app/[org]/[dataset]/page.tsx @@ -5,6 +5,10 @@ export default function DatasetRootPage({ }: { params: { org: string; dataset: string }; }) { - redirect(`/${params.org}/${params.dataset}/episode_0`); - return null; + const episodeN = process.env.EPISODES + ?.split(/\s+/) + .map((x) => parseInt(x.trim(), 10)) + .filter((x) => !isNaN(x))[0] ?? 0; + + redirect(`/${params.org}/${params.dataset}/episode_${episodeN}`); } diff --git a/lerobot/html_dataset_visualizer/src/app/page.tsx b/lerobot/html_dataset_visualizer/src/app/page.tsx index 8403b3f60..b4029cf36 100644 --- a/lerobot/html_dataset_visualizer/src/app/page.tsx +++ b/lerobot/html_dataset_visualizer/src/app/page.tsx @@ -9,6 +9,15 @@ export default function Home({ }: { searchParams: { [key: string]: string | undefined }; }) { + // Redirect to the first episode of the dataset if REPO_ID is defined + if (process.env.REPO_ID) { + const episodeN = process.env.EPISODES + ?.split(/\s+/) + .map((x) => parseInt(x.trim(), 10)) + .filter((x) => !isNaN(x))[0] ?? 0; + + redirect(`/${process.env.REPO_ID}/episode_${episodeN}`); + } // sync with hf.co/spaces URL params if (searchParams.path) { redirect(searchParams.path); diff --git a/lerobot/scripts/visualize_dataset_html.py b/lerobot/scripts/visualize_dataset_html.py index d0c8f1ace..13f7545dd 100644 --- a/lerobot/scripts/visualize_dataset_html.py +++ b/lerobot/scripts/visualize_dataset_html.py @@ -53,80 +53,33 @@ python lerobot/scripts/visualize_dataset_html.py \ """ import argparse -import csv import json -import logging -import re -import shutil -import tempfile -from io import StringIO +import os from pathlib import Path +import subprocess +import atexit +import signal +import sys +import logging -import numpy as np -import pandas as pd -import requests -from flask import Flask, redirect, render_template, request, url_for +from flask import Flask, jsonify, redirect, send_file, url_for -from lerobot import available_datasets from lerobot.common.datasets.lerobot_dataset import LeRobotDataset -from lerobot.common.datasets.utils import IterableNamespace +from lerobot.common.datasets.utils import INFO_PATH, DEFAULT_PARQUET_PATH, DEFAULT_VIDEO_PATH from lerobot.common.utils.utils import init_logging -def run_server( - dataset: LeRobotDataset | IterableNamespace | None, - episodes: list[int] | None, +def run_data_server( + dataset: LeRobotDataset | None, host: str, - port: str, - static_folder: Path, - template_folder: Path, -): - app = Flask(__name__, static_folder=static_folder.resolve(), template_folder=template_folder.resolve()) - app.config["SEND_FILE_MAX_AGE_DEFAULT"] = 0 # specifying not to cache + port: int, +) -> Path | None: + init_logging() - @app.route("/") - def hommepage(dataset=dataset): - if dataset: - dataset_namespace, dataset_name = dataset.repo_id.split("/") - return redirect( - url_for( - "show_episode", - dataset_namespace=dataset_namespace, - dataset_name=dataset_name, - episode_id=0, - ) - ) + data_server = Flask(__name__) + data_server.config["SEND_FILE_MAX_AGE_DEFAULT"] = 0 # specifying not to cache - dataset_param, episode_param = None, None - all_params = request.args - if "dataset" in all_params: - dataset_param = all_params["dataset"] - if "episode" in all_params: - episode_param = int(all_params["episode"]) - - if dataset_param: - dataset_namespace, dataset_name = dataset_param.split("/") - return redirect( - url_for( - "show_episode", - dataset_namespace=dataset_namespace, - dataset_name=dataset_name, - episode_id=episode_param if episode_param is not None else 0, - ) - ) - - featured_datasets = [ - "lerobot/aloha_static_cups_open", - "lerobot/columbia_cairlab_pusht_real", - "lerobot/taco_play", - ] - return render_template( - "visualize_dataset_homepage.html", - featured_datasets=featured_datasets, - lerobot_datasets=available_datasets, - ) - - @app.route("//") + @data_server.route("//") def show_first_episode(dataset_namespace, dataset_name): first_episode_id = 0 return redirect( @@ -138,256 +91,133 @@ def run_server( ) ) - @app.route("///episode_") - def show_episode(dataset_namespace, dataset_name, episode_id, dataset=dataset, episodes=episodes): - repo_id = f"{dataset_namespace}/{dataset_name}" + @data_server.route("///resolve/main/meta/info.json") + def serve_info_json(dataset_namespace, dataset_name): try: - if dataset is None: - dataset = get_dataset_info(repo_id) + return send_file(dataset.root / INFO_PATH, mimetype="application/json") except FileNotFoundError: - return ( - "Make sure to convert your LeRobotDataset to v2 & above. See how to convert your dataset at https://github.com/huggingface/lerobot/pull/461", - 400, - ) - dataset_version = ( - str(dataset.meta._version) if isinstance(dataset, LeRobotDataset) else dataset.codebase_version - ) - match = re.search(r"v(\d+)\.", dataset_version) - if match: - major_version = int(match.group(1)) - if major_version < 2: - return "Make sure to convert your LeRobotDataset to v2 & above." + return jsonify({"error": "File not found"}), 404 + except Exception as e: + return jsonify({"error": f"Server error: {str(e)}"}), 500 - 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 - if isinstance(dataset, LeRobotDataset) - else dataset.total_frames, - "num_episodes": dataset.num_episodes - if isinstance(dataset, LeRobotDataset) - else dataset.total_episodes, - "fps": dataset.fps, - } - if isinstance(dataset, LeRobotDataset): - video_paths = [ - dataset.meta.get_video_file_path(episode_id, key) for key in dataset.meta.video_keys - ] - videos_info = [ - { - "url": url_for("static", filename=str(video_path).replace("\\", "/")), - "filename": video_path.parent.name, - } - for video_path in video_paths - ] - tasks = dataset.meta.episodes[episode_id]["tasks"] - else: - video_keys = [key for key, ft in dataset.features.items() if ft["dtype"] == "video"] - videos_info = [ - { - "url": f"https://huggingface.co/datasets/{repo_id}/resolve/main/" - + dataset.video_path.format( - episode_chunk=int(episode_id) // dataset.chunks_size, - video_key=video_key, - episode_index=episode_id, - ), - "filename": video_key, - } - for video_key in video_keys - ] - - response = requests.get( - f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/episodes.jsonl", timeout=5 - ) - response.raise_for_status() - # Split into lines and parse each line as JSON - tasks_jsonl = [json.loads(line) for line in response.text.splitlines() if line.strip()] - - filtered_tasks_jsonl = [row for row in tasks_jsonl if row["episode_index"] == episode_id] - tasks = filtered_tasks_jsonl[0]["tasks"] - - videos_info[0]["language_instruction"] = tasks - - if episodes is None: - episodes = list( - range(dataset.num_episodes if isinstance(dataset, LeRobotDataset) else dataset.total_episodes) - ) - - return render_template( - "visualize_dataset_template.html", - episode_id=episode_id, - episodes=episodes, - dataset_info=dataset_info, - videos_info=videos_info, - episode_data_csv_str=episode_data_csv_str, - columns=columns, - ignored_columns=ignored_columns, - ) - - app.run(host=host, port=port) - - -def get_ep_csv_fname(episode_id: int): - ep_csv_fname = f"episode_{episode_id}.csv" - return ep_csv_fname - - -def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index): - """Get a csv str containing timeseries data of an episode (e.g. state and action). - This file will be loaded by Dygraph javascript to plot data in real time.""" - columns = [] - - selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] in ["float32", "int32"]] - 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"] - - for column_name in selected_columns: - dim_state = ( - dataset.meta.shapes[column_name][0] - if isinstance(dataset, LeRobotDataset) - else dataset.features[column_name].shape[0] - ) - - if "names" in dataset.features[column_name] and dataset.features[column_name]["names"]: - column_names = dataset.features[column_name]["names"] - while not isinstance(column_names, list): - column_names = list(column_names.values())[0] - else: - column_names = [f"{column_name}_{i}" for i in range(dim_state)] - columns.append({"key": column_name, "value": column_names}) - - header += column_names - - selected_columns.insert(0, "timestamp") - - if isinstance(dataset, LeRobotDataset): - from_idx = dataset.episode_data_index["from"][episode_index] - to_idx = dataset.episode_data_index["to"][episode_index] - data = ( - dataset.hf_dataset.select(range(from_idx, to_idx)) - .select_columns(selected_columns) - .with_format("pandas") - ) - else: - repo_id = dataset.repo_id - - url = f"https://huggingface.co/datasets/{repo_id}/resolve/main/" + dataset.data_path.format( - episode_chunk=int(episode_index) // dataset.chunks_size, episode_index=episode_index - ) - df = pd.read_parquet(url) - data = df[selected_columns] # Select specific columns - - rows = np.hstack( - ( - np.expand_dims(data["timestamp"], axis=1), - *[np.vstack(data[col]) for col in selected_columns[1:]], - ) - ).tolist() - - # Convert data to CSV string - csv_buffer = StringIO() - csv_writer = csv.writer(csv_buffer) - # Write header - csv_writer.writerow(header) - # Write data rows - csv_writer.writerows(rows) - csv_string = csv_buffer.getvalue() - - return csv_string, columns, ignored_columns - - -def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]: - # get first frame of episode (hack to get video_path of the episode) - first_frame_idx = dataset.episode_data_index["from"][ep_index].item() - return [ - dataset.hf_dataset.select_columns(key)[first_frame_idx][key]["path"] - for key in dataset.meta.video_keys - ] - - -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.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 get_dataset_info(repo_id: str) -> IterableNamespace: - response = requests.get( - f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json", timeout=5 + @data_server.route( + "///resolve/main/data/chunk-/episode_.parquet" ) - response.raise_for_status() # Raises an HTTPError for bad responses - dataset_info = response.json() - dataset_info["repo_id"] = repo_id - return IterableNamespace(dataset_info) + def serve_parquet_file(dataset_namespace, dataset_name, episode_chunk, episode_index): + try: + # Format the path with the captured parameters + file_path = DEFAULT_PARQUET_PATH.format(episode_chunk=episode_chunk, episode_index=episode_index) + full_path = dataset.root / file_path -def visualize_dataset_html( - dataset: LeRobotDataset | None, - episodes: list[int] | None = None, - output_dir: Path | None = None, - serve: bool = True, - host: str = "127.0.0.1", - port: int = 9090, - force_override: bool = False, -) -> Path | None: - init_logging() + return send_file(full_path, mimetype="application/octet-stream") + except FileNotFoundError: + return jsonify({"error": "File not found"}), 404 + except Exception as e: + return jsonify({"error": f"Server error: {str(e)}"}), 500 - template_dir = Path(__file__).resolve().parent.parent / "templates" - - if output_dir is None: - # Create a temporary directory that will be automatically cleaned up - output_dir = tempfile.mkdtemp(prefix="lerobot_visualize_dataset_") - - output_dir = Path(output_dir) - if output_dir.exists(): - if force_override: - shutil.rmtree(output_dir) - else: - logging.info(f"Output directory already exists. Loading from it: '{output_dir}'") - - output_dir.mkdir(parents=True, exist_ok=True) - - static_dir = output_dir / "static" - static_dir.mkdir(parents=True, exist_ok=True) - - if dataset is None: - if serve: - run_server( - dataset=None, - episodes=None, - host=host, - port=port, - static_folder=static_dir, - template_folder=template_dir, + @data_server.route( + "///resolve/main/videos/chunk-//episode_.mp4" + ) + def serve_video_file(dataset_namespace, dataset_name, episode_chunk, video_key, episode_index): + try: + # Format the path with the captured parameters + file_path = DEFAULT_VIDEO_PATH.format( + episode_chunk=episode_chunk, video_key=video_key, episode_index=episode_index ) - else: - # Create a simlink from the dataset video folder containing mp4 files to the output directory - # so that the http server can get access to the mp4 files. - if isinstance(dataset, LeRobotDataset): - ln_videos_dir = static_dir / "videos" - if not ln_videos_dir.exists(): - ln_videos_dir.symlink_to((dataset.root / "videos").resolve().as_posix()) - if serve: - run_server(dataset, episodes, host, port, static_dir, template_dir) + # Assuming 'dataset' object has a 'root' attribute + full_path = dataset.root / file_path + + return send_file(full_path, mimetype="video/mp4") + except FileNotFoundError: + return jsonify({"error": "Video file not found"}), 404 + except Exception as e: + return jsonify({"error": f"Server error: {str(e)}"}), 500 + + log = logging.getLogger("werkzeug") + log.setLevel(logging.ERROR) + + data_server.run(host=host, port=get_local_data_server_port(port)) + + +def is_npm_available(): + try: + subprocess.run(["npm", "--version"], capture_output=True, text=True, check=True) + return True + except (subprocess.CalledProcessError, FileNotFoundError): + return False + + +def build_react_app(script_dir: Path): + next_dir = script_dir.parent / "html_dataset_visualizer" + next_build_dir = next_dir / ".next" + if not next_build_dir.exists() or not next_build_dir.is_dir(): + print("Building React.js app ...") + subprocess.run(["npm", "ci"], cwd=next_dir) + subprocess.run(["npm", "run", "build"], cwd=next_dir) + + package_json_path = next_dir / "package.json" + build_id_path = next_build_dir / "BUILD_ID" + with open(package_json_path, "r") as f: + package_data = json.load(f) + package_version = package_data.get("version", "") + with open(build_id_path, "r") as f: + build_id = f.read().strip() + + if package_version != build_id: + print("Building React.js app ...") + subprocess.run(["npm", "ci"], cwd=next_dir) + subprocess.run(["npm", "run", "build"], cwd=next_dir) + + +def run_react_app( + repo_id: str, + script_dir: Path, + load_from_hf_hub: bool, + host: str, + port: int, + episodes: list[int] | None = None, +): + next_dir = script_dir.parent / "html_dataset_visualizer" + + env = os.environ.copy() + env["REPO_ID"] = repo_id + if not load_from_hf_hub: + env["DATASET_URL"] = f"http://{host}:{get_local_data_server_port(port)}" + if episodes: + env["EPISODES"] = " ".join(map(str, episodes)) + + process = subprocess.Popen( + ["npm", "run", "start", "--", f"--port={port}"], cwd=next_dir, env=env, preexec_fn=os.setsid + ) + + def cleanup(): + if process.poll() is None: # Process still running + print("Cleaning up React server...") + try: + os.killpg(os.getpgid(process.pid), signal.SIGTERM) + process.wait(timeout=5) + except (ProcessLookupError, subprocess.TimeoutExpired): + # Force kill if graceful termination fails + try: + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + except ProcessLookupError: + pass + + def signal_handler(sig, frame): + cleanup() + sys.exit(0) + + signal.signal(signal.SIGINT, signal_handler) + atexit.register(cleanup) # Also cleanup on normal exit + + return process + + +def get_local_data_server_port(port: str): + """Returns the port used by the local data server.""" + return str(int(port) + 1) def main(): @@ -396,8 +226,8 @@ def main(): parser.add_argument( "--repo-id", type=str, - default=None, help="Name of hugging face repositery containing a LeRobotDataset dataset (e.g. `lerobot/pusht` for https://huggingface.co/datasets/lerobot/pusht).", + required=True, ) parser.add_argument( "--root", @@ -418,18 +248,6 @@ def main(): default=None, help="Episode indices to visualize (e.g. `0 1 5 6` to load episodes of index 0, 1, 5 and 6). By default loads all episodes.", ) - parser.add_argument( - "--output-dir", - type=Path, - default=None, - help="Directory path to write html files and kickoff a web server. By default write them to 'outputs/visualize_dataset/REPO_ID'.", - ) - parser.add_argument( - "--serve", - type=int, - default=1, - help="Launch web server.", - ) parser.add_argument( "--host", type=str, @@ -442,13 +260,6 @@ def main(): default=9090, help="Web port used by the http server.", ) - parser.add_argument( - "--force-override", - type=int, - default=0, - help="Delete the output directory if it exists already.", - ) - parser.add_argument( "--tolerance-s", type=float, @@ -466,16 +277,21 @@ def main(): load_from_hf_hub = kwargs.pop("load_from_hf_hub") root = kwargs.pop("root") tolerance_s = kwargs.pop("tolerance_s") + host = kwargs.pop("host") + port = kwargs.pop("port") + episodes = kwargs.pop("episodes") - dataset = None - if repo_id: - dataset = ( - LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s) - if not load_from_hf_hub - else get_dataset_info(repo_id) - ) + if not is_npm_available(): + raise RuntimeError("npm is not available. Please install it to use this script.") - visualize_dataset_html(dataset, **vars(args)) + script_dir = Path(__file__).parent.absolute() + + build_react_app(script_dir) + run_react_app(repo_id, script_dir, load_from_hf_hub, host, port, episodes) + + if not load_from_hf_hub: + dataset = LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s) + run_data_server(dataset, host, port) if __name__ == "__main__":