Spaces:
Running
Running
| """HDF5 Export — 把批次构建为单一 .h5 数据集。 | |
| 与 LeRobot Export 节点同源数据(视频/手部骨骼/传感器/标注),由 | |
| app.hdf5_export.build_hdf5_dataset 完成:视频逐帧抽帧 → uint8 数组 | |
| (gzip 压缩),表列按 LeRobot observation 语义组织成层级结构。 | |
| """ | |
| from app.processing import ProcessingModule, JobContext, ArtifactRef, field | |
| from app.processing.registry import register | |
| from app.processing.theme import EXPORT_COLOR | |
| class HDF5ExportModule(ProcessingModule): | |
| slug = "hdf5_export" | |
| version = "1.1" | |
| category = "export" | |
| label = "HDF5 Export" | |
| icon = "ant-design:database-outlined" | |
| color = EXPORT_COLOR | |
| inputs = ({"key": "data", "label": "Exportable Data"},) | |
| outputs = ({"key": "dataset", "label": "Dataset"},) | |
| default_config = {"compression": "gzip", "level": 4} | |
| config_schema = ( | |
| field("compression", "select", "Compression", "gzip", options=["gzip", "lzf"]), | |
| field("level", "number", "Compression level", 4, min=0, max=9), | |
| ) | |
| execution_target = "worker" | |
| capabilities = ("dataset_export", "hdf5") | |
| 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.hdf5_export import build_hdf5_dataset | |
| compression = str(ctx.config.get("compression", "gzip")) | |
| level = int(ctx.config.get("level", 4) or 4) | |
| # 连接驱动:输入卡片(视频)连到本节点 → 抽帧该视频;没连 → 不抽帧。 | |
| include_video_keys = [ | |
| r.source_key for r in ctx.incoming.values() | |
| if r.kind == "video" and r.source_key | |
| ] | |
| 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 | |
| ] | |
| 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 | |
| ] | |
| 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) | |
| out_path = build_hdf5_dataset( | |
| episode_id, [episode_id], ctx.output_root, | |
| include_video_keys=include_video_keys, | |
| hand_keypoints_paths=hand_keypoints_paths, | |
| hand_3d_paths=hand_3d_paths, | |
| hand_3d_unit=hand_3d_unit, | |
| compression=compression, level=level) | |
| if not out_path.exists(): | |
| ctx.skip("HDF5 build produced no file") | |
| return { | |
| "dataset": ctx.ref("dataset", out_path, metadata={ | |
| "root": str(out_path.relative_to(ctx.output_root)) if out_path != ctx.output_root else ".", | |
| "episode_id": episode_id, | |
| }), | |
| } | |