Dataset Merger Tool README

一、简介

merge_lerobot_dataset是一个功能强大的Python脚本，专门用于合并多个具有相似结构的数据集。它特别适用于包含视频、状态和动作数据的lerobot自采数据集，能够高效地整合分散的episodes和分散的任务资源，为后续的模型训练等工作提供便利。

二、功能概述

数据集合并：将多个源数据集文件夹合并为一个统一的数据集，简化数据管理。
索引重编号：重新编号所有的episode索引和任务索引，确保合并后数据的连续性和一致性。
向量维度填充：自动检测并填充向量维度，使所有数据在observation.state和action等方面具有一致性。
统计信息合并：智能合并多个数据集的统计信息，正确处理复杂的数据结构，如图像特征的嵌套结构。
图像视频验证：如果有图片 images 文件夹的话，自动检测视频 / 图像 / 元数据文件之间的数量关系。
视频文件处理：正确处理和复制视频文件，保持视频与其他数据之间的正确索引关系，支持多种视频存储结构。
元数据更新：更新所有元数据文件，准确反映合并后的数据集结构。
图像文件处理：复制图像文件，保持图片与其他数据之间的正确索引关系。

三、安装

本脚本依赖于以下Python库：

numpy
pandas

安装依赖项的命令如下：

pip install numpy pandas

四、使用方法

（一）基本用法

通过命令行运行工具，示例如下：

python dataset_merger.py --sources /path/to/dataset1 /path/to/dataset2 /path/to/dataset3 --output /path/to/output_dataset

（二）命令行参数

--sources：源数据集文件夹路径列表，至少需要指定一个源数据集路径。
--output：输出数据集文件夹路径，用于指定合并后数据集的存储位置。
--state_max_dim：状态向量的最大维度，默认值为32。
--action_max_dim：动作向量的最大维度，默认值为32。
--fps：数据集的帧率，默认值为20。
--copy_images: 是否将图像从源文件夹复制且合并到输出文件夹。(default: False)

（三）示例

python dataset_merger.py --sources ./robot_dataset_1 ./robot_dataset_2 --output ./merged_dataset --state_max_dim 32 --action_max_dim 18 --fps 30

五、数据集格式

此工具假设输入数据集具有以下结构：

dataset/
├── meta/
│   ├── episodes.jsonl       # 包含每个episode的元数据
│   ├── episodes_stats.jsonl # 每个episode的统计数据
│   ├── info.json           # 数据集全局信息
│   ├── stats.json          # 全局统计信息
│   └── tasks.jsonl         # 任务定义
├── data/                   # 包含parquet格式的episode数据
│   └── chunk-xxx/
│       └── episode_xxxxxx.parquet
├── images/                   # 可选的图片文件
│   └── episode_xxxxxx/
│       └── frame_xxxxxx.png
└── videos/                 # 可选的视频文件
    └── chunk-xxx/
        └── video_key/
            └── episode_xxxxxx.mp4

六、功能细节

数据一致性处理：自动检测并填充状态和动作向量维度，确保所有数据具有一致的维度，满足机器学习算法对数据格式的要求。
索引管理：重新编号所有episode和任务索引，同时维护帧索引的连续性，避免数据混乱。
统计信息合并：智能合并多个数据集的统计数据，能够正确处理复杂的数据结构，如图像特征的嵌套结构，确保统计信息的准确性和完整性。
视频文件处理：正确复制视频文件，并保持视频与其他数据之间的正确索引关系，支持多种视频存储结构，保证视频数据与其他数据的同步性。
任务映射：自动检测并合并相同的任务描述，创建新的任务索引映射，方便对任务进行统一管理和调用。
数据预验证：在执行合并操作前对数据集执行全面的预验证，检查视频帧数、图片数量与元数据中记录的帧长度的一致性，确保合并后数据的准确性和完整性，并自动修复可修复的问题（从图片重新编码视频）。
图像文件管理：复制和整理图像文件，保持正确的命名规则和目录结构，确保图像与视频和其他数据保持正确的索引对应关系，支持按需启用图像复制功能来优化存储空间使用。

七、注意事项

确保所有源数据集具有兼容的结构，否则可能导致合并失败或数据错误。
合并后的数据集可能占用较大磁盘空间，在进行合并操作前，请确保有足够的存储空间。
对于非常大的数据集，合并过程可能需要较长时间，请耐心等待。
图像文件夹占用磁盘空间很大，默认不开启 copy_images 参数。本工具在处理 images 文件夹下的图片时，仅支持 PNG 格式的图片，且要求图片文件名为 frame_XXXXXX.png（X为6位数字，例如 frame_000001.png）。合并过程中会自动检测并处理这些 PNG 图片。

八、常见问题

Q: 合并不同维度的数据集会发生什么？ A: 工具会自动检测最大维度并用零填充较小维度的向量，确保所有数据具有一致的维度。
Q: 如何处理不同FPS的数据集？ A: 暂时只支持相同FPS的数据集合并。
Q: 能否只合并某些特定episode？ A: 当前版本会合并所有数据。如需更精细的控制，您可以先筛选数据集，然后再进行合并。也可以全部合并，然后使用lerobot加载特定的episode。