iDP3复现代码数据预处理全流程（二）——vis_dataset.py
vis_dataset.py 主要作用在于点云数据的可视化，并可以做一些简单的预处理
关键参数基本都在 vis_dataset.sh 中定义了，需要改动的仅以下两点：
1. 点云图像保存位置，因为 dataset_path 被设置为了绝对路径，因此需要相应修改：
save_dir = f"{dataset_path}/{episode_idx}" # 设置当前集的保存目录2. 点云视频保存位置，对应修改：
if vis_cloud:
# 将图像序列转换为视频
os.system(f"ffmpeg -r 10 -i {save_dir}/%d.png -vcodec mpeg4 -y visualizations/{dataset_path}/{episode_idx}.mp4")运行生成：
更详细解释如下：
目录
1 库函数调用
2 创建 ArgumentParser 对象，解析命令行参数
3 打开指定路径 Zarr 数据集并获取数据
4 分割数据、处理并保存
5 回放图像和点云数据
6 将点云图像存为视频
1 库函数调用
import zarr # 导入 zarr 库，用于处理 zarr 格式的数据
import cv2 # 导入 OpenCV 库，用于图像处理
from termcolor import cprint # 从 termcolor 库导入 cprint 函数，用于彩色打印输出
import time # 导入 time 库，用于时间相关操作
from tqdm import tqdm # 导入 tqdm 库，用于显示进度条
import visualizer # 导入自定义的 visualizer 模块，用于点云可视化
import os # 导入 os 库，用于操作系统相关功能
import argparse # 导入 argparse 库，用于解析命令行参数
import numpy as np # 导入 numpy 库，用于数值计算除了 visualizer 为自定义库函数，其余均为标准库
2 创建 ArgumentParser 对象，解析命令行参数
# 创建 ArgumentParser 对象，用于解析命令行参数
parser = argparse.ArgumentParser()
parser.add_argument("--dataset_path", type=str, default="data/box_zarr") # 添加 dataset_path 参数，默认值为 "data/box_zarr"
parser.add_argument("--use_img", type=int, default=0) # 添加 use_img 参数，默认值为 0（不使用图像数据）
parser.add_argument("--vis_cloud", type=int, default=0) # 添加 vis_cloud 参数，默认值为 0（不可视化点云数据）
parser.add_argument("--use_pc_color", type=int, default=0) # 添加 use_pc_color 参数，默认值为 0（不使用点云颜色）
parser.add_argument("--downsample", type=int, default=0) # 添加 downsample 参数，默认值为 0（不下采样）
# 解析命令行参数
args = parser.parse_args()
use_img = args.use_img
dataset_path = args.dataset_path
vis_cloud = args.vis_cloud
use_pc_color = args.use_pc_color
downsample = args.downsample命令行参数均在 vis_dataset.sh 中定义，若未定义，则会使用默认参数
3 打开指定路径 Zarr 数据集并获取数据
# 使用 zarr 打开指定路径的数据集
with zarr.open(dataset_path) as zf:
print(zf.tree()) # 打印数据集的树状结构
# 获取数据
if use_img:
all_img = zf['data/img'] # 获取图像数据
all_point_cloud = zf['data/point_cloud'] # 获取点云数据
all_episode_ends = zf['meta/episode_ends'] # 获取集结束标记打开 dataset_path 位置数据集，并获取对应数据
4 分割数据、处理并保存
# 根据 episode_ends 划分每一集的数据
for episode_idx, episode_end in enumerate(all_episode_ends):
if episode_idx == 0:
if use_img:
img_episode = all_img[:episode_end] # 获取第一集的图像数据
point_cloud_episode = all_point_cloud[:episode_end] # 获取第一集的点云数据
else:
if use_img:
img_episode = all_img[all_episode_ends[episode_idx-1]:episode_end] # 获取当前集的图像数据
point_cloud_episode = all_point_cloud[all_episode_ends[episode_idx-1]:episode_end] # 获取当前集的点云数据
save_dir = f"visualizations/{dataset_path}/{episode_idx}" # 设置当前集的保存目录
if vis_cloud:
os.makedirs(save_dir, exist_ok=True) # 创建保存目录（如果需要可视化点云）
cprint(f"replay episode {episode_idx}", "green") # 打印当前集的重放信息，使用绿色字体5 回放图像和点云数据
# 回放每一帧的数据
for i in range(point_cloud_episode.shape[0]):
pc = point_cloud_episode[i] # 获取当前帧的点云数据
# 如果需要下采样
if downsample:
num_points = 4096 # 设置下采样点数
idx = np.random.choice(pc.shape[0], num_points, replace=False) # 随机选择点
pc = pc[idx] # 获取下采样后的点云数据
if use_img:
img = img_episode[i] # 获取当前帧的图像数据
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 将图像从 BGR 转换为 RGB
cv2.imshow('img', img) # 显示图像
cv2.waitKey(1) # 等待 1 毫秒
time.sleep(0.05) # 暂停 0.05 秒
if vis_cloud:
if not use_pc_color:
pc = pc[:, :3] # 如果不使用点云颜色，只保留前三个维度（XYZ）
visualizer.visualize_pointcloud(pc, img_path=f"{save_dir}/{i}.png") # 可视化点云并保存图像
print(f"vis cloud saved to {save_dir}/{i}.png") # 打印保存路径
print(f"frame {i}/{point_cloud_episode.shape[0]}") # 打印当前帧的处理进度6 将点云图像存为视频
if vis_cloud:
# 将图像序列转换为视频
os.system(f"ffmpeg -r 10 -i {save_dir}/%d.png -vcodec mpeg4 -y {dataset_path}/{episode_idx}.mp4")使用 ffmpeg 命令将保存在{save_dir}目录下的一系列PNG图片（按序号命名）转换为帧率为10帧每秒的MPEG-4格式视频
并保存到 {dataset_path}/{episode_idx}.mp4，如果输出文件已存在，则会直接覆盖