Spaces:
Running
Running
| """LeRobot Export — 把批次构建为标准 LeRobot v3 数据集。 | |
| 真正导出由 app.lerobot_export.build_lerobot_dataset 完成(视频/深度/ | |
| 传感器/标注 → meta/info.json + 分块 parquet)。本模块把该实现接入 | |
| 工作流节点:跑完工作流,数据集先生成到 worker 临时目录，完成后交给系统级 | |
| 导出缓存；不在项目根目录生成 processed/，也不把导出副本混入源数据集。 | |
| """ | |
| from app.processing import ProcessingModule, JobContext, ArtifactRef, field | |
| from app.processing.registry import register | |
| from app.processing.theme import EXPORT_COLOR | |
| class LeRobotExportModule(ProcessingModule): | |
| slug = "lerobot_export" | |
| version = "1.0" | |
| category = "export" | |
| label = "LeRobot Export" | |
| icon = "ant-design:cloud-server-outlined" | |
| color = EXPORT_COLOR | |
| inputs = ({"key": "data", "label": "Exportable Data"},) | |
| outputs = ({"key": "dataset", "label": "Dataset"},) | |
| # 默认排除 UMI 触觉力矩阵:每帧 250×250×3(187500 个整数),实测占 | |
| # canonical parquet 体积的九成以上,且 99.99% 是零 —— 训练用不到, | |
| # 却让单集导出从 2.5 秒涨到 57 秒(写盘 + 逐元素统计 + 结果包上传), | |
| # 结果 zip 也从几 MB 涨到几十 MB。要保留就把它从配置里删掉。 | |
| _DEFAULT_EXCLUDE = ",".join(( | |
| "observation.gripper_left_force_matrix", | |
| "observation.gripper_right_force_matrix", | |
| )) | |
| default_config = {"version": "v3.0", "split_ratio": 0.9, | |
| "shard_size": 100000, "exclude_columns": _DEFAULT_EXCLUDE} | |
| config_schema = ( | |
| field("version", "select", "Version", "v3.0", options=["v2.1", "v3.0"]), | |
| field("exclude_columns", "string", "Exclude columns (comma separated)", | |
| _DEFAULT_EXCLUDE), | |
| ) | |
| execution_target = "worker" | |
| capabilities = ("dataset_export", "lerobot") | |
| def run(self, ctx: JobContext) -> dict[str, ArtifactRef] | None: | |
| if not ctx.incoming: | |
| ctx.skip("No upstream data — skipped") | |
| episode_id = str(ctx.job.get("episode_id") or "") | |
| if not episode_id: | |
| ctx.skip("No episode_id in job") | |
| from app.lerobot_export import build_lerobot_dataset | |
| split_ratio = float(ctx.config.get("split_ratio", 0.9)) | |
| # 布局版本:config_schema 的 version 选项(v2.1/v3.0)直接驱动 | |
| # 导出版本由卡片配置决定；源项目始终保持本地 canonical v2.1 | |
| # 三目录结构，v3.0 只在系统级导出缓存中生成。 | |
| version = str(ctx.config.get("version") or "v3.0") | |
| # 导出范围由工作流连接决定:输入卡片(视频)连到本节点 → 数据集 | |
| # 包含该视频;没连的相机 → 默认不导出视频(传感器/标注仍全量)。 | |
| include_video_keys = [ | |
| r.source_key for r in ctx.incoming.values() | |
| if r.kind == "video" and r.source_key | |
| ] | |
| # 手部骨骼产物:处理节点(mediapipe_hand)连到本节点 → 一并写入 | |
| # 数据集(连接驱动);没连 → 不写入。 | |
| # 注意必须 resolve:ref.path 是相对 output_root 的,直接传给 | |
| # build_lerobot_dataset 会因工作目录不同而读不到文件。 | |
| hand_keypoints_paths = [ | |
| str(p) for r in ctx.incoming.values() | |
| if r.kind == "hand_keypoints" and r.path | |
| for p in (ctx.resolve(r),) if p | |
| ] | |
| # 双目三角化产物(stereo_triangulate 连到本节点)→ hand_*_world | |
| # 变为米制 3D 坐标(优先于 2D 归一化坐标)。 | |
| hand_3d_paths = [ | |
| str(p) for r in ctx.incoming.values() | |
| if r.kind == "hand_3d" and r.path | |
| for p in (ctx.resolve(r),) if p | |
| ] | |
| # 右目(辅助视角)骨骼产物(stereo_triangulate 的 hand_3d#right | |
| # "#" 兄弟键,沿边自动透传)连到本节点 → 以 _rcam 后缀列组写入 | |
| # 数据集(与主数据并存)。 | |
| hand_3d_right_paths = [ | |
| str(p) for r in ctx.incoming.values() | |
| if r.kind == "hand_3d#right" and r.path | |
| for p in (ctx.resolve(r),) if p | |
| ] | |
| # hand_3d 产物的 unit 语义(manifest 声明,如深度图抬升的 camera_meters) | |
| # → 透传给数据集 feature 声明(缺省保持 mediapipe_world_relative)。 | |
| hand_3d_unit = next( | |
| (r.metadata.get("unit") for r in ctx.incoming.values() | |
| if r.kind == "hand_3d" and getattr(r, "metadata", None) | |
| and r.metadata.get("unit")), | |
| None) | |
| # 数据集先输出到 worker 临时目录，API 完成阶段会把导出产品放到 | |
| # state/exports，不在项目根目录新增 processed/。 | |
| # 排除列:配置里逗号分隔。留空 = 全量导出(旧行为)。 | |
| exclude_columns = [ | |
| name.strip() for name | |
| in str(ctx.config.get("exclude_columns", "") or "").split(",") | |
| if name.strip() | |
| ] | |
| dataset_dir = build_lerobot_dataset(episode_id, [episode_id], | |
| ctx.output_root, split_ratio, | |
| include_video_keys=include_video_keys, | |
| hand_keypoints_paths=hand_keypoints_paths, | |
| hand_3d_paths=hand_3d_paths, | |
| hand_3d_right_paths=hand_3d_right_paths, | |
| version=version, | |
| hand_3d_unit=hand_3d_unit, | |
| exclude_columns=exclude_columns) | |
| info = dataset_dir / "meta" / "info.json" | |
| if not info.exists(): | |
| ctx.skip("Dataset build produced no meta/info.json") | |
| return { | |
| "dataset": ctx.ref("dataset", info, metadata={ | |
| "root": str(dataset_dir.relative_to(ctx.output_root)) if dataset_dir != ctx.output_root else ".", | |
| "episode_id": episode_id, | |
| }), | |
| } | |