iDP3复现代码数据预处理全流程（二）——vis_dataset.py

vis_dataset.py 主要作用在于点云数据的可视化，并可以做一些简单的预处理

关键参数基本都在 vis_dataset.sh 中定义了，需要改动的仅以下两点：

点云图像保存位置，因为 dataset_path 被设置为了绝对路径，因此需要相应修改：

python 复制代码

save_dir = f"{dataset_path}/{episode_idx}"  # 设置当前集的保存目录

点云视频保存位置，对应修改：

python 复制代码

        if vis_cloud:
            # 将图像序列转换为视频
            os.system(f"ffmpeg -r 10 -i {save_dir}/%d.png -vcodec mpeg4 -y visualizations/{dataset_path}/{episode_idx}.mp4")

运行生成：

更详细解释如下：

[1 库函数调用](#1 库函数调用)

[2 创建 ArgumentParser 对象，解析命令行参数](#2 创建 ArgumentParser 对象，解析命令行参数)

[3 打开指定路径 Zarr 数据集并获取数据](#3 打开指定路径 Zarr 数据集并获取数据)

[4 分割数据、处理并保存](#4 分割数据、处理并保存)

[5 回放图像和点云数据](#5 回放图像和点云数据)

[6 将点云图像存为视频](#6 将点云图像存为视频)

1 库函数调用

python 复制代码

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 对象，解析命令行参数

python 复制代码

# 创建 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 数据集并获取数据

python 复制代码

# 使用 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 分割数据、处理并保存

python 复制代码

    # 根据 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 回放图像和点云数据

python 复制代码

        # 回放每一帧的数据
        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 将点云图像存为视频

python 复制代码

        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，如果输出文件已存在，则会直接覆盖