Spaces:
Running
Running
File size: 3,091 Bytes
d646dd7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """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
@register
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,
}),
}
|