forked from tangger/lerobot
[visualization] Ignore 2d or 3d data for now (#809)
This commit is contained in:
@@ -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]:
|
||||
|
||||
Reference in New Issue
Block a user