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

server/2024/12/31 22:30:02/

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

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

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

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

2. 点云视频保存位置,对应修改:

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

运行生成:

更详细解释如下:

目录

1 库函数调用

2 创建 ArgumentParser 对象,解析命令行参数

3 打开指定路径 Zarr 数据集并获取数据

4 分割数据、处理并保存

5 回放图像和点云数据

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")  # 打印当前集的重放信息,使用绿色字体

回放图像和点云数据

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 转换为 RGBcv2.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,如果输出文件已存在,则会直接覆盖


http://www.ppmy.cn/server/154275.html

相关文章

动态规划简介:爱的初体验

动态规划是一个让人欲仙欲死、死去活来的科目,茫然无措时让你痛不欲生,柳暗花明时让你喜不自胜。 小伙,不能啊! 你要心怀有爱,你要用全身心的爱来迎接今天,迎接动态规划。 一、动态规划的概念 动态规划…

不安全物联网的轻量级加密:综述

Abstract 本文综述了针对物联网(IoT)的轻量级加密解决方案。这项综述全面覆盖了从轻量级加密方案到不同类型分组密码的比较等多个方面。同时,还对硬件与软件解决方案之间的比较进行了讨论,并分析了当前最受信赖且研究最深入的分组…

Python 敲电子木鱼,见机甲佛祖,修赛博真经

Python 敲电子木鱼,见机甲佛祖,修赛博真经 相关资源文件已经打包成EXE文件,可下载相关资源压缩包后双击直接运行程序,且文章末尾已附上相关源码,以供大家学习交流,博主主页还有更多Python相关程序案例&…

【AI日记】24.12.28 kaggle 比赛 2-16

【AI论文解读】【AI知识点】【AI小项目】【AI战略思考】【AI日记】 工作 参加:kaggle 比赛 Regression with an Insurance Dataset时间:7 小时 读书 书名:教育的本质时间:1.5 小时 律己 工作时间:差作息&#xff…

STM32-笔记11-手写带操作系统的延时函数

1、为什么带操作系统的延时函数,和笔记10上的延时函数不能使用同一种? 因为笔记10的延时函数在每次调用的时候,会一直开关定时器,而在FreeRTOS操作系统中,SysTick定时器当作时基使用。 时基是一个时间显示的基本单位。…

机器人历史

机器人的历史可以追溯到古代。古希腊神话中,有一位名叫赫菲斯托斯的神,据说他制造了可移动的机械人。另外,古代中国、古印度和古埃及也有关于机械人的传说和故事。 然而,现代机器人的发展始于20世纪。1921年,捷克作家…

大数据-261 实时数仓 - 业务数据库表结构 交易订单、订单产品、产品分类、商家店铺、地域组织表

点一下关注吧!!!非常感谢!!持续更新!!! Java篇开始了! MyBatis 更新完毕目前开始更新 Spring,一起深入浅出! 目前已经更新到了: H…

操作系统实验三 存储管理

实验三 存储管理 一、实验目的 通过实验使学生了解可变式分区管理使用的主要数据结构,分配、回收的主要技术,了解最优适应分配、最坏适应分配、最先适应分配和循环适应分配等分配算法。基本能达到下列具体的目标: 掌握初步进程在内存中的映…