Open X-Embodiment数据集下载及数据解析（可视化成视频及对应的instruction）_open x-embodiment 数据集下载

Open X-Embodiment数据集下载及数据解析（可视化成视频及对应的instruction）

• 1 数据下载
• • 1.1 官方的Colab往后翻可以找到每个数据对应的路径:
  • 1.2 在服务器上下载数据
• 2 数据可视化
• • 2.1 环境
  • 2.2 代码
  • 2.3 可视化结果

1 数据下载

1.1 官方的Colab往后翻可以找到每个数据对应的路径:

数据集项目链接
官方的Colab

gs://gresearch/robotics/fractal20220817_data/0.1.0 has size 111.07 GiBgs://gresearch/robotics/kuka/0.1.0 has size 778.02 GiBgs://gresearch/robotics/bridge/0.1.0 has size 387.49 GiBgs://gresearch/robotics/taco_play/0.1.0 has size 47.77 GiBgs://gresearch/robotics/jaco_play/0.1.0 has size 9.24 GiBgs://gresearch/robotics/berkeley_cable_routing/0.1.0 has size 4.67 GiBgs://gresearch/robotics/roboturk/0.1.0 has size 45.39 GiBgs://gresearch/robotics/nyu_door_opening_surprising_effectiveness/0.1.0 has size 7.12 GiBgs://gresearch/robotics/viola/0.1.0 has size 10.40 GiBgs://gresearch/robotics/berkeley_autolab_ur5/0.1.0 has size 76.39 GiBgs://gresearch/robotics/toto/0.1.0 has size 127.66 GiBgs://gresearch/robotics/language_table/0.0.1 has size 399.23 GiBgs://gresearch/robotics/columbia_cairlab_pusht_real/0.1.0 has size 2.80 GiBgs://gresearch/robotics/stanford_kuka_multimodal_dataset_converted_externally_to_rlds/0.1.0 has size 31.98 GiBgs://gresearch/robotics/nyu_rot_dataset_converted_externally_to_rlds/0.1.0 has size 5.33 MiBgs://gresearch/robotics/stanford_hydra_dataset_converted_externally_to_rlds/0.1.0 has size 72.48 GiBgs://gresearch/robotics/austin_buds_dataset_converted_externally_to_rlds/0.1.0 has size 1.49 GiBgs://gresearch/robotics/nyu_franka_play_dataset_converted_externally_to_rlds/0.1.0 has size 5.18 GiBgs://gresearch/robotics/maniskill_dataset_converted_externally_to_rlds/0.1.0 has size 151.05 GiBgs://gresearch/robotics/cmu_franka_exploration_dataset_converted_externally_to_rlds/0.1.0 has size 602.24 MiBgs://gresearch/robotics/ucsd_kitchen_dataset_converted_externally_to_rlds/0.1.0 has size 1.33 GiBgs://gresearch/robotics/ucsd_pick_and_place_dataset_converted_externally_to_rlds/0.1.0 has size 3.53 GiBgs://gresearch/robotics/austin_sailor_dataset_converted_externally_to_rlds/0.1.0 has size 18.85 GiBgs://gresearch/robotics/austin_sirius_dataset_converted_externally_to_rlds/0.1.0 has size 6.55 GiBgs://gresearch/robotics/bc_z/0.1.0 has size 80.54 GiBgs://gresearch/robotics/usc_cloth_sim_converted_externally_to_rlds/0.1.0 has size 254.52 MiBgs://gresearch/robotics/utokyo_pr2_opening_fridge_converted_externally_to_rlds/0.1.0 has size 360.57 MiBgs://gresearch/robotics/utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds/0.1.0 has size 829.37 MiBgs://gresearch/robotics/utokyo_saytap_converted_externally_to_rlds/0.1.0 has size 55.34 MiBgs://gresearch/robotics/utokyo_xarm_pick_and_place_converted_externally_to_rlds/0.1.0 has size 1.29 GiBgs://gresearch/robotics/utokyo_xarm_bimanual_converted_externally_to_rlds/0.1.0 has size 138.44 MiBgs://gresearch/robotics/robo_net/1.0.0 has size 799.91 GiBgs://gresearch/robotics/berkeley_mvp_converted_externally_to_rlds/0.1.0 has size 12.34 GiBgs://gresearch/robotics/berkeley_rpt_converted_externally_to_rlds/0.1.0 has size 40.64 GiBgs://gresearch/robotics/kaist_nonprehensile_converted_externally_to_rlds/0.1.0 has size 11.71 GiBgs://gresearch/robotics/stanford_mask_vit_converted_externally_to_rlds/0.1.0 has size 76.17 GiBgs://gresearch/robotics/tokyo_u_lsmo_converted_externally_to_rlds/0.1.0 has size 335.71 MiBgs://gresearch/robotics/dlr_sara_pour_converted_externally_to_rlds/0.1.0 has size 2.92 GiBgs://gresearch/robotics/dlr_sara_grid_clamp_converted_externally_to_rlds/0.1.0 has size 1.65 GiBgs://gresearch/robotics/dlr_edan_shared_control_converted_externally_to_rlds/0.1.0 has size 3.09 GiBgs://gresearch/robotics/asu_table_top_converted_externally_to_rlds/0.1.0 has size 737.60 MiBgs://gresearch/robotics/stanford_robocook_converted_externally_to_rlds/0.1.0 has size 124.62 GiBgs://gresearch/robotics/eth_agent_affordances/0.1.0 has size 17.27 GiBgs://gresearch/robotics/imperialcollege_sawyer_wrist_cam/0.1.0 has size 81.87 MiBgs://gresearch/robotics/iamlab_cmu_pickup_insert_converted_externally_to_rlds/0.1.0 has size 50.29 GiBgs://gresearch/robotics/uiuc_d3field/0.1.0 has size 15.82 GiBgs://gresearch/robotics/utaustin_mutex/0.1.0 has size 20.79 GiBgs://gresearch/robotics/berkeley_fanuc_manipulation/0.1.0 has size 8.85 GiBgs://gresearch/robotics/cmu_play_fusion/0.1.0 has size 6.68 GiBgs://gresearch/robotics/cmu_stretch/0.1.0 has size 728.06 MiBgs://gresearch/robotics/berkeley_gnm_recon/0.1.0 has size 18.73 GiBgs://gresearch/robotics/berkeley_gnm_cory_hall/0.1.0 has size 1.39 GiBgs://gresearch/robotics/berkeley_gnm_sac_son/0.1.0 has size 7.00 GiB

1.2 在服务器上下载数据

下载gsutil库

pip install gsutil

下载数据集

gsutil -m cp -r gs://gresearch/robotics/fractal20220817_data/0.1.0[具体的数据集路径根据需求替换成上面的] /path/to/your/dir

或者

/snap/bin/gsutil -m cp -r gs://gresearch/robotics/fractal20220817_data/0.1.0[具体的数据集路径根据需求替换成上面的] /path/to/your/dir

找不到的数据集可以先查看是否存在：

/snap/bin/gsutil ls gs://gresearch/robotics/jaco_play[替换成数据集的Registered Dataset Name]

具体的Registered Dataset Name可以在数据信息表中找到

2 数据可视化

2.1 环境

pip install tensorflow-datasets==4.9.3

2.2 代码

import numpy as npimport tensorflow_datasets as tfdsfrom PIL import Imagefrom IPython import displayimport osdisplay_key = \'image\'datasets_name = \"bridge\"b = tfds.builder_from_directory(f\"/path/to/Open_X_Embodiment_Datasets/{datasets_name}\")ds = b.as_dataset(split=\'train\') # 具体可以根据需求改output_dir = f\'{datasets_name}_videos\'os.makedirs(output_dir, exist_ok=True)instructions_file_path = os.path.join(output_dir, f\'{datasets_name}.txt\')state_file_path = os.path.join(output_dir, f\'state.txt\')# 遍历数据集for idx, episode in enumerate(ds): # 为每个视频创建一个文件夹 video_folder = os.path.join(output_dir, f\'video_{idx}\') os.makedirs(video_folder, exist_ok=True) # 提取该视频的所有帧 frames = episode[\'steps\'] # 遍历每一帧并保存 state_list = [] for frame_idx, step in enumerate(frames): # 每个数据集image的特征名字不一样，具体要看数据集下载好后的 features.json 文件中对应的字段是什么 image = step[\'observation\'][image] # fractal20220817_data # image = step[\'observation\'][\"agentview_rgb\"] # viola # image = step[\'observation\'][\"image\"] # bridge # 获取自然语言指令，具体要看数据集下载好后的 features.json 文件对应的字段是什么 # natural_language_instruction = step[\"language_instruction\"].numpy().decode(\'utf-8\') # for ucsd、berkeley_fanuc_manipulation natural_language_instruction = step[\'observation\'][\"natural_language_instruction\"].numpy().decode(\'utf-8\') state_list.append(step[\'observation\'][\"state\"]) # 将图像转换为 PIL 格式 image_pil = Image.fromarray(image.numpy()) # 保存图像，文件名格式为 frame_{frame_idx}.png output_path = os.path.join(video_folder, f\'frame_{frame_idx}.png\') image_pil.save(output_path) with open(state_file_path, \'a\') as f: f.write(f\"state {idx}: {state_list}\\n\") with open(instructions_file_path, \'a\') as f: f.write(f\"Video {idx} Instruction: {natural_language_instruction}\\n\") print(f\"第 {idx} 个视频的所有帧已保存到: {video_folder}, 该视频共有{frame_idx + 1}帧\")print(\"所有视频的帧提取完成。\")

• 代码主要参考数据集官方的Colab
• 不同的数据集要根据 features.json 文件中的字段来修改代码中的字段，比如 fractal20220817_data 数据集的 features.json 文件中图片的字段是 image，所以代码中：image = step[\'observation\'][image]
  
• 自然语言指令的字段是 natural_language_instruction，所以代码中：natural_language_instruction = step[\'observation\'][\"natural_language_instruction\"].numpy().decode(\'utf-8\')
  
• 如果想提取其他信息也是同理，主要根据 features.json 中的字段对上面的代码进行调整

2.3 可视化结果

代码运行完毕后会保存图片帧：

• 以及对应的 instruction，在.txt文件中：
  
• 最后，如果要将图片转为视频，可以参考以下代码：

import cv2import osimport redef natural_sort_key(s): # 提取字符串中的数字部分，并将其转换为整数 return [int(text) if text.isdigit() else text for text in re.split(\'([0-9]+)\', s)]def images_to_video(image_folder, video_name, fps=10): # 获取文件夹中的所有图像文件 # images = [img for img in os.listdir(image_folder) if img.endswith(\".jpg\") or img.endswith(\".png\")] filenames = sorted(os.listdir(image_folder), key=natural_sort_key) # images.sort() # 确保图像按顺序排序 # 确定视频的宽度和高度 first_image = cv2.imread(os.path.join(image_folder, filenames[0])) height, width, layers = first_image.shape # 创建视频写入对象 fourcc = cv2.VideoWriter_fourcc(*\'mp4v\') # 可以选择不同的编码格式 video = cv2.VideoWriter(video_name, fourcc, fps, (width, height)) for image in filenames: img_path = os.path.join(image_folder, image) frame = cv2.imread(img_path) video.write(frame) # 将每一帧写入视频 video.release() # 释放视频写入对象 cv2.destroyAllWindows() # 关闭所有 OpenCV 窗口 print(f\"视频 \'{video_name}\' 已成功创建！\")# 使用示例image_folder = \"/path/to/Open_X_Embodiment_Datasets/RT-1_Robot_Action_videos/video_0\" # 替换为图像文件夹路径video_name = \"/path/to/output/RT-1_videos_277.mp4\" # 输出视频文件名images_to_video(image_folder, video_name, fps=5)