"""Local-file LeRobot export (v2.1 / v3.0). The review application keeps the uploaded episode untouched. This module creates a derived LeRobot dataset with standard tabular/video metadata. Depth is carried by the per-camera raw depth-code video (``is_depth_map``) plus the per-frame ``observation.depth..valid`` marker; the original 16-bit PNG sequence is not copied into the dataset. ``version`` 决定布局(与官方结构一一对应,两者互不兼容): - v2.1:每 episode 一个文件 —— ``data/chunk-{episode_chunk:03d}/ episode_{episode_index:06d}.parquet``、``videos/{video_key}/chunk- {episode_chunk:03d}/episode_{episode_index:06d}.mp4``;episode 元数据写在官方旧版 ``meta/episodes.jsonl``,并附带 ``meta/episodes_stats.jsonl``;info.json 用 v2.1 路径模板并含 ``total_videos``/``total_chunks``。 - v3.0(默认):合并分片 —— ``data/chunk-{n}/file-{i:03d}.parquet``(单文件 含全部 episode)、``videos/chunk-{n}/{video_key}/file-{i:03d}.mp4``; episode 元数据在官方的 ``meta/episodes/chunk-{n}/file-{i:03d}.parquet``。 """ from __future__ import annotations import json import re import shutil from pathlib import Path from typing import Any from app.localstore import get_episode, list_annotations from app.config import settings from app.lerobot_v21 import ( DEPTH_MAX_MM, DEPTH_MIN_MM, DEPTH_QMAX, DEPTH_QP, DEPTH_VIDEO_ENCODING, depth_video_encoder_args, quantize_depth, ) def _json_default(value: Any): if hasattr(value, "tolist"): return value.tolist() if hasattr(value, "item"): return value.item() raise TypeError(f"Unsupported value: {type(value)!r}") def _write_json(path: Path, value: Any) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text( json.dumps(value, ensure_ascii=False, indent=2, default=_json_default), encoding="utf-8", ) def _write_tasks_parquet(output_dir: Path, task_descriptions: dict[int, str]) -> None: """Write the v3 task index in the native Parquet representation. ``tasks.jsonl`` is kept as a compatibility companion, but v3 consumers expect ``meta/tasks.parquet``. 任务文本必须存成 DataFrame index: 官方加载器按 ``tasks.iloc[task_idx].name`` 取任务名(样本行里的 ``task`` 字段),存普通列会返回整数下标而非任务文本(实测)。 """ import pandas as pd items = sorted(task_descriptions.items()) frame = pd.DataFrame( {"task_index": [int(task_id) for task_id, _ in items]}, index=pd.Index([str(description) for _, description in items], name="task"), ) path = output_dir / "meta" / "tasks.parquet" path.parent.mkdir(parents=True, exist_ok=True) frame.to_parquet(path) def _quantiles(arr: "np.ndarray") -> list[float]: import numpy as np return [float(q) for q in np.nanpercentile(arr, [1, 10, 50, 90, 99])] def _nanpercentile_axis0(mat: "np.ndarray", qs) -> "np.ndarray": """等价于 ``np.nanpercentile(mat, qs, axis=0)``,但整矩阵一次算完。 numpy 对带 axis 的 nanpercentile 会退化成 ``apply_along_axis`` —— 逐列 调 Python。触觉力矩阵是 (帧数, 187500),那就是每次 18.75 万次 1-D 调用, 实测 4.5 秒;两个矩阵就是 9 秒。这里改成"整矩阵排序一次 + 按各列有效值 个数做向量化索引",全走 C 层,实测 0.30 秒(15 倍),数值逐元素一致 (linear 插值口径与 numpy 相同,已对无 NaN/含 NaN/单行/整列 NaN/常数列 逐一比对)。 NaN 排到每列末尾后不影响取值 —— 有效值个数 counts 把它们排除在外。 """ import numpy as np qs = np.asarray(qs, dtype=np.float64) valid = np.isfinite(mat) counts = valid.sum(axis=0) safe = np.where(valid, mat, np.inf) safe.sort(axis=0) n = counts[None, :] # linear 插值:virtual_index = q/100 * (n - 1),在 lo/hi 之间按小数部分插值 pos = (n - 1) * (qs[:, None] / 100.0) lo = np.floor(pos) frac = pos - lo hi = np.ceil(pos) last = np.maximum(n - 1, 0) lo_i = np.clip(lo, 0, last).astype(np.intp) hi_i = np.clip(hi, 0, last).astype(np.intp) cols = np.arange(mat.shape[1])[None, :] out = safe[lo_i, cols] * (1 - frac) + safe[hi_i, cols] * frac out[:, counts == 0] = np.nan # 与 nanpercentile 一致:整列无效 → NaN return out def _write_stats(output_dir: Path, video_features: dict[str, dict] | None = None, data_table=None) -> None: """LeRobot 官方 meta/stats.json:每个数值列 mean/std/min/max(训练归一化用)。 官方数据集(如 aloha_mobile)在 meta/stats.json 提供统计量,训练管线 读它做 observation/action 归一化。覆盖三类特征: - 标量列:mean/std/min/max + q01-q99 + count; - 定长数值向量列(如 hand_left_world 63 维):逐维统计(官方 pose.head 同款,扁平列表); - 视频特征:抽帧解码,逐通道统计(归一化到 [0,1],官方 observation.images.* 同款 [3,1,1] 嵌套)。 变长嵌套数组列与字符串列不做数值统计。 ``data_table``: 调用方刚写进 data/ 的 arrow table。传了就直接转 pandas, 省掉把整个 data 目录从磁盘再读一遍;不传(旧调用点/表不可用)则退回 原磁盘路径,行为不变。 """ import numpy as np import pandas as pd try: if data_table is not None: # ``pd.read_parquet`` 内部就是 ``pq.read_table(...).to_pandas()``, # 所以这里转出来的 DataFrame 与重读磁盘逐列同 dtype、同值(实测 # 校验过),但省掉一次全量往返 —— data/ 实测约 13MB/集,50 集 # 就是 650MB 从 SSHFS 重新搬进内存再 concat。 df = data_table.to_pandas() else: parquet_files = sorted((output_dir / "data").rglob("*.parquet")) if not parquet_files: return frames = [pd.read_parquet(p) for p in parquet_files] if not frames: return df = pd.concat(frames, ignore_index=True) except Exception: return stats: dict[str, dict] = {} for col in df.columns: s = df[col] try: if not pd.api.types.is_numeric_dtype(s): continue arr = s.to_numpy() if arr.ndim != 1: continue if any(isinstance(v, (list, tuple)) for v in arr): continue # 变长嵌套数组列跳过 except Exception: continue try: num = pd.to_numeric(s, errors="coerce").to_numpy(dtype=float) except Exception: continue if not np.isfinite(num).any(): continue q = _quantiles(num) # 只写数值统计量。不要再加 dtype/shape 之类的元数据字段:官方 # 加载器会把 stats 里的每一项都转成张量做归一化,字符串字段直接 # 崩在 torch.from_numpy(np.ndarray of type numpy.str_)。官方数据集 # 的 stats 也只有 min/max/mean/std/count。 stats[col] = { "mean": float(np.nanmean(num)), "std": float(np.nanstd(num)), "min": float(np.nanmin(num)), "max": float(np.nanmax(num)), "q01": q[0], "q10": q[1], "q50": q[2], "q90": q[3], "q99": q[4], "count": int((~np.isnan(num)).sum()), } # 定长数值向量列(手部世界坐标 63 维等):逐维统计,官方 pose.head 格式 try: for col in df.columns: s = df[col] if not isinstance(s.dtype, np.dtype) or s.dtype != object: continue vals = s.dropna().tolist() # pandas 把 pyarrow list 列还原成 ndarray(而非 list),两种都接受 if not vals or not all( isinstance(v, (list, tuple, np.ndarray)) for v in vals): continue lengths = {len(v) for v in vals} if len(lengths) != 1: continue try: mat = np.array(vals, dtype=float) except Exception: continue if mat.ndim != 2 or not np.isfinite(mat).any(): continue q = _nanpercentile_axis0(mat, [1, 10, 50, 90, 99]) stats[col] = { "min": np.nanmin(mat, axis=0).tolist(), "max": np.nanmax(mat, axis=0).tolist(), "mean": np.nanmean(mat, axis=0).tolist(), "std": np.nanstd(mat, axis=0).tolist(), "q01": q[0].tolist(), "q10": q[1].tolist(), "q50": q[2].tolist(), "q90": q[3].tolist(), "q99": q[4].tolist(), "count": int(np.isfinite(mat).all(axis=1).sum()), } except Exception: pass # 视频特征:RGB 抽帧逐通道统计(归一化 [0,1],官方 # observation.images.* 格式)。metric depth 使用 pyav/ffmpeg 读出 # gray12le 后反量化为毫米;它不是普通 uint8 视频,不能走 OpenCV # 解码后统一除以 255 的 RGB 路径。 try: import cv2 for vkey, meta in (video_features or {}).items(): if meta.get("dtype") != "video": continue videos_root = output_dir / "videos" videos = sorted( path for path in videos_root.rglob("*.mp4") if vkey in path.relative_to(videos_root).parts ) if not videos: continue shape = meta.get("shape") or [0, 0, 3] channels = int(shape[2]) if len(shape) > 2 else 3 video_info = meta.get("video_info") or meta.get("info") or {} is_metric_depth = bool( video_info.get("video.is_depth_map") and str(video_info.get("video.pix_fmt") or "").lower() in {"gray12le", "gray16le", "gray10le"} ) # OpenCV converts high-bit-depth gray12le to an 8-bit display # image. That loses the metric depth codes and produces stats in # an unrelated [0, 1] range after /255. Use the project's depth # reader instead, which returns physical millimetres. if is_metric_depth: from app.lerobot_v21 import DepthVideoReader samples: list = [] for vpath in videos: try: reader = DepthVideoReader(vpath) except Exception as exc: print(f"[lerobot_export] depth stats skipped for {vpath}: {exc}") continue total = int(reader.frame_count or 0) step = max(1, total // 200) if total > 200 else 1 read_idx = 0 try: while True: depth_mm = reader.read_mm() if depth_mm is None: break if read_idx % step == 0: # Keep the same one-channel shape used by # the RGB statistics below. depth_mm = depth_mm.astype(np.float32, copy=False) # LeRobot's depth dequantizer clamps decoded # values to [depth_min, depth_max]. The raw # reader preserves the transport sentinel 0, # so apply the same rule before calculating # q01/q99; otherwise q01=0 while the GUI sees # those pixels as depth_min (usually 100 mm), # compressing the visible range near yellow. depth_min_m = float(video_info.get("video.depth_min") or 0.1) depth_max_m = float(video_info.get("video.depth_max") or 5.0) depth_min_mm = depth_min_m * 1000.0 depth_max_mm = depth_max_m * 1000.0 depth_mm = np.nan_to_num( depth_mm, nan=depth_min_mm, posinf=depth_max_mm, neginf=depth_min_mm, ) depth_mm = np.clip(depth_mm, depth_min_mm, depth_max_mm) samples.append(depth_mm[..., None]) read_idx += 1 finally: reader.close() if not samples: continue stack = np.stack(samples) # (T, H, W, 1), values in mm per_channel = stack.reshape(-1, stack.shape[-1]) def _nested(vals): return [[[float(v)] for v in vals]] q = np.percentile(per_channel, [1, 10, 50, 90, 99], axis=0) stats[vkey] = { "min": _nested(per_channel.min(axis=0)), "max": _nested(per_channel.max(axis=0)), "mean": _nested(per_channel.mean(axis=0)), "std": _nested(per_channel.std(axis=0)), "q01": _nested(q[0]), "q10": _nested(q[1]), "q50": _nested(q[2]), "q90": _nested(q[3]), "q99": _nested(q[4]), "count": [int(per_channel.shape[0])], } continue # 增量累积:全部集的采样帧若一次性留在内存里会爆 —— 960×1280 # 的一帧 float32 是 14.75MB,34 集 × 约 200 采样帧 ≈ 100GB。 # 这里只保留跑动统计量(min/max/和/平方和/计数)与一个固定大小的 # 分位数蓄水池,内存与集数无关。 ch_sum = np.zeros(channels, dtype=np.float64) ch_sq = np.zeros(channels, dtype=np.float64) ch_min = np.full(channels, np.inf, dtype=np.float64) ch_max = np.full(channels, -np.inf, dtype=np.float64) pixel_count = 0 # 蓄水池上限与每帧贡献:分位数只需近似,但样本必须**跨帧、跨集** # 均匀铺开 —— 只取前几帧会让分位数偏向最早的那一集。每帧按固定 # 步长抽约 1200 个像素,800 万像素上限可覆盖数千帧。 reservoir_cap = 8_000_000 per_frame_budget = 1200 taken = 0 reservoir: list = [] for vpath in videos: cap = cv2.VideoCapture(str(vpath)) if not cap.isOpened(): continue total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 step = max(1, total // 200) if total > 200 else 1 # 顺序读、每 step 帧采样一次(随机 seek 在 MP4/HEVC 上 # 每次都要回溯关键帧,长视频统计会慢一个数量级) read_idx = 0 while True: ok, frame = cap.read() if not ok: break if read_idx % step == 0 and frame.ndim == 3: # 先按步长抽样再转 float:一帧 960×1280 有 120 万像素, # 对整帧做 astype(float32)/255 是这一步的主要开销, # 而统计只需要千级样本。抽样后再运算快约三个数量级。 flat_u8 = frame[..., :channels].reshape(-1, channels) stride_u8 = max(1, flat_u8.shape[0] // per_frame_budget) flat = (flat_u8[::stride_u8].astype(np.float32) / 255.0) ch_sum += flat.sum(axis=0, dtype=np.float64) ch_sq += np.square(flat, dtype=np.float64).sum(axis=0) ch_min = np.minimum(ch_min, flat.min(axis=0)) ch_max = np.maximum(ch_max, flat.max(axis=0)) pixel_count += flat.shape[0] if taken < reservoir_cap: # flat 已是抽样结果,直接入池 take = min(reservoir_cap - taken, flat.shape[0]) reservoir.append(flat[:take].copy()) taken += take read_idx += 1 cap.release() if not pixel_count: continue def _nested(vals): return [[[float(v)] for v in vals]] mean = ch_sum / pixel_count var = np.maximum(ch_sq / pixel_count - mean * mean, 0.0) pooled = np.concatenate(reservoir, axis=0) if reservoir else None if pooled is not None and pooled.shape[0]: q = np.percentile(pooled, [1, 10, 50, 90, 99], axis=0) else: q = np.stack([mean] * 5) stats[vkey] = { "min": _nested(ch_min), "max": _nested(ch_max), "mean": _nested(mean), "std": _nested(np.sqrt(var)), "q01": _nested(q[0]), "q10": _nested(q[1]), "q50": _nested(q[2]), "q90": _nested(q[3]), "q99": _nested(q[4]), "count": [int(pixel_count)], } except Exception: pass if stats: _write_json(output_dir / "meta" / "stats.json", stats) def _write_episodes_stats_v21(output_dir: Path, episode_row_ranges: list[tuple], rows: list[dict], episode_video_meta_records: list, video_features: dict[str, dict] | None) -> None: """v2.1 legacy ``meta/episodes_stats.jsonl``:每集一行 {episode_index, stats}。 官方 v2.1 每集一份归一化统计(转换脚本/训练管线读取):标量列 min/max/mean/std/count、定长向量列逐维统计、视频抽帧逐通道统计 (归一化 [0,1])。count 必须数组格式(官方校验要求)。v3.0 已移除此 文件(统计收敛进 meta/stats.json + episodes parquet)。 """ import numpy as np import cv2 def _nested(vals): # 官方转换脚本要求逐通道统计形状 (C,1,1)(实测 (1,C,1) 报 # "Shape of quantile 'min' must be (3,1,1) or (1,1,1)") return [[[float(v)]] for v in vals] lines: list[str] = [] for (ep_index, _task_id, _length, start, end), vmeta in zip( episode_row_ranges, episode_video_meta_records): ep_rows = rows[start:end] if not ep_rows: continue stats: dict[str, dict] = {} for key in ep_rows[0].keys(): vals = [row.get(key) for row in ep_rows] if all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in vals): num = np.array(vals, dtype=float) if not np.isfinite(num).any(): continue stats[key] = { "min": float(np.nanmin(num)), "max": float(np.nanmax(num)), "mean": float(np.nanmean(num)), "std": float(np.nanstd(num)), "count": [len(vals)], } elif all(isinstance(v, (list, tuple, np.ndarray)) for v in vals): try: mat = np.array([list(v) for v in vals], dtype=float) except Exception: continue if mat.ndim == 2 and np.isfinite(mat).any(): stats[key] = { "min": np.nanmin(mat, axis=0).tolist(), "max": np.nanmax(mat, axis=0).tolist(), "mean": np.nanmean(mat, axis=0).tolist(), "std": np.nanstd(mat, axis=0).tolist(), "count": [len(vals)], } # 该集每路视频一个文件(v2.1 布局),抽帧逐通道统计 chunk = ep_index // 1000 for vkey, meta in (video_features or {}).items(): if meta.get("dtype") != "video": continue vpath = (output_dir / "videos" / vkey / f"chunk-{chunk:03d}" / f"episode_{ep_index:06d}.mp4") if not vpath.exists(): continue shape = meta.get("shape") or [0, 0, 3] channels = int(shape[2]) if len(shape) > 2 else 3 cap = cv2.VideoCapture(str(vpath)) if not cap.isOpened(): continue samples: list = [] total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 step = max(1, total // 200) if total > 200 else 1 read_idx = 0 while True: ok, frame = cap.read() if not ok: break if read_idx % step == 0 and frame.ndim == 3: samples.append(frame.astype(np.float32)[..., :channels] / 255.0) read_idx += 1 cap.release() if not samples: continue stack = np.stack(samples) per_channel = stack.reshape(-1, stack.shape[-1]) stats[vkey] = { "min": _nested(per_channel.min(axis=0)), "max": _nested(per_channel.max(axis=0)), "mean": _nested(per_channel.mean(axis=0)), "std": _nested(per_channel.std(axis=0)), "count": [len(samples)], } lines.append(json.dumps({"episode_index": ep_index, "stats": stats}, ensure_ascii=False)) if lines: (output_dir / "meta" / "episodes_stats.jsonl").write_text( "\n".join(lines) + "\n", encoding="utf-8") def _key_matches_any(candidate: str, keys: list[str]) -> bool: """候选相机名与任一 key 双向子串匹配(与 find_videos 同规则)。 输入卡片 config.source_key(如 head_left_rgb)与批次相机名可能 不完全相等,子串匹配保证连接关系能对上实际视频。 """ cl = candidate.lower() return any(k.lower() in cl or cl in k.lower() for k in keys) def _source_path(path_value: str | None) -> Path | None: if not path_value: return None path = Path(path_value) if path.is_absolute(): return path return settings.storage_root / path def _episode_storage_index(session_dir: Path, episode_id: str | None = None) -> int | None: if not episode_id: return None try: from app.project_dataset import episode_row row = episode_row(Path(session_dir), str(episode_id)) return int(row["episode_index"]) if row and row.get("episode_index") is not None else None except (TypeError, ValueError, OSError): return None def _depth_root(session_dir: Path, episode_id: str | None = None) -> Path: root = Path(session_dir) / "depth" if root.is_dir(): return root namespaced = Path(session_dir) / "meta" / "depth" / str(episode_id or "") if episode_id and namespaced.is_dir(): return namespaced return Path(session_dir) / "meta" / "depth" def _depth_video_paths(session_dir: Path, episode_id: str | None = None) -> list[tuple[str, Path]]: """Discover every uploaded depth video without assuming a device name. Depth streams are stored separately from ``camera_streams`` under ``depth//``. They are therefore invisible to the normal RGB video loop below. Keep the source name from the directory (or filename for a flat layout) so different depth devices use the same export path. """ # Canonical LeRobot v2.1 batches keep depth video beside RGB streams. # These streams are included here as already-encoded inputs. PNG depth # sequences are not part of the active format. result: list[tuple[str, Path]] = [] seen: set[str] = set() try: from app.lerobot_v21 import is_depth_source, iter_video_streams episode_index = _episode_storage_index(Path(session_dir), episode_id) for source, path in iter_video_streams(Path(session_dir) / "videos"): if (is_depth_source(source) and source not in seen and (episode_index is None or path.stem == f"episode_{episode_index:06d}")): seen.add(source) result.append((source, path)) except Exception: pass return result def _probe_video(path: Path) -> tuple[int, float, int, int]: try: import cv2 cap = cv2.VideoCapture(str(path)) count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) fps = float(cap.get(cv2.CAP_PROP_FPS) or 0) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) cap.release() return count, fps, width, height except Exception: return 0, 0.0, 0, 0 def _probe_depth_video(path: Path) -> dict: """Probe a depth stream with ffprobe first, OpenCV as fallback. OpenCV builds without 12-bit HEVC support silently report zeros for gray12le streams. ffprobe returns the true pixel format and dimensions; the container duration estimates the frame count when ``nb_frames`` is absent from the stream entry (common for HEVC). """ import json import subprocess result = {"frames": 0, "fps": 0.0, "width": 0, "height": 0, "pix_fmt": ""} try: completed = subprocess.run( ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=codec_name,pix_fmt,width,height,r_frame_rate,nb_frames:" "format=duration", "-of", "json", str(path)], capture_output=True, text=True, timeout=15, check=False, ) payload = json.loads(completed.stdout or "{}") streams = payload.get("streams") or [] if streams: stream = streams[0] result["pix_fmt"] = str(stream.get("pix_fmt") or "") for name in ("width", "height"): try: result[name] = int(stream.get(name) or 0) except (TypeError, ValueError): pass try: result["frames"] = int(stream.get("nb_frames") or 0) except (TypeError, ValueError): pass rate = str(stream.get("r_frame_rate") or "").split("/") try: if len(rate) == 2 and float(rate[1]) > 0: result["fps"] = float(rate[0]) / float(rate[1]) except (TypeError, ValueError, ZeroDivisionError): pass if not result["frames"] and result["fps"] > 0: try: duration = float((payload.get("format") or {}).get("duration") or 0) except (TypeError, ValueError): duration = 0.0 if duration > 0: result["frames"] = int(round(duration * result["fps"])) except (OSError, subprocess.SubprocessError, ValueError, TypeError, json.JSONDecodeError): pass count, fps, width, height = _probe_video(path) # OpenCV fallback result["frames"] = result["frames"] or count result["fps"] = result["fps"] or float(fps or 0.0) result["width"] = result["width"] or int(width or 0) result["height"] = result["height"] or int(height or 0) return result def _copy_depth_video_canonical(source: Path, destination: Path) -> None: """Copy a depth stream and normalize HEVC's MP4 sample entry. Some uploaded gray12le streams use ``hev1`` while the project contract uses ``hvc1``. This is a container-only remux (``-c copy``): depth codes are not decoded, re-quantized, or re-encoded. Non-gray12le legacy visualization streams are copied unchanged. """ import subprocess try: probe = subprocess.run( ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=codec_name,pix_fmt,codec_tag_string", "-of", "json", str(source)], capture_output=True, text=True, check=True, ) stream = (json.loads(probe.stdout).get("streams") or [{}])[0] is_metric_hevc = (stream.get("codec_name") == "hevc" and stream.get("pix_fmt") == "gray12le") already_canonical = stream.get("codec_tag_string") == "hvc1" except (OSError, subprocess.SubprocessError, ValueError, TypeError, json.JSONDecodeError): is_metric_hevc = False already_canonical = False if not is_metric_hevc or already_canonical: shutil.copy2(source, destination) return tmp = destination.with_suffix(".remux.mp4") tmp.unlink(missing_ok=True) try: subprocess.run( ["ffmpeg", "-y", "-loglevel", "error", "-i", str(source), "-map", "0:v:0", "-c:v", "copy", "-tag:v", "hvc1", "-movflags", "+faststart", str(tmp)], check=True, capture_output=True, text=True, ) if not tmp.exists() or tmp.stat().st_size == 0: raise RuntimeError("empty remux output") tmp.replace(destination) except (OSError, subprocess.SubprocessError, RuntimeError) as exc: tmp.unlink(missing_ok=True) # Keep export usable for installations without a compatible ffmpeg; # metadata still records the actual gray12le stream. print(f"[lerobot_export] depth hvc1 remux skipped: {exc}") shutil.copy2(source, destination) def _write_depth_video(dest: Path, pngs: list[Path], frame_count: int, fps: float, depth_scale: float | None, depth_min: float = 0.01, depth_max: float = 10.0, shift: float = 3.5, use_log: bool = True) -> tuple[int, float, int, int]: """PNG 深度序列(uint16 毫米)→ canonical metric depth-code 视频。 深度值本身存进视频(替代 JET 伪彩):每帧 uint16 → 米 → 裁剪 → log 量化 12bit 码 → gray12le → libx265 CQP qp=6。 下游按 canonical 元数据反量化回 metric 深度。缺帧写全 0 码 (无效值,配合 observation.depth..valid 列识别)。 """ import cv2 import numpy as np import subprocess import shutil exe = shutil.which("ffmpeg") if not exe or not pngs: return 0, 0.0, 0, 0 # 帧号 → PNG 文件名:0-based(000000.png 存在)或 1-based(000001.png) offset = 0 if (pngs[0].parent / "000000.png").exists() else 1 by_stem = {p.stem: p for p in pngs} sample = cv2.imread(str(pngs[0]), cv2.IMREAD_UNCHANGED) if sample is None: return 0, 0.0, 0, 0 height, width = sample.shape[:2] scale = float(depth_scale) if depth_scale else 0.001 # 默认毫米 def _frame_to_codes(frame_index: int) -> np.ndarray: path = by_stem.get(f"{frame_index + offset:06d}") if path is None: return np.zeros((height, width), dtype=np.uint16) arr = cv2.imread(str(path), cv2.IMREAD_UNCHANGED) if arr is None: return np.zeros((height, width), dtype=np.uint16) depth_mm = arr.astype(np.float64) * scale * 1000.0 return quantize_depth(depth_mm) dest.parent.mkdir(parents=True, exist_ok=True) tmp_dest = dest.with_suffix(".tmp.mp4") tmp_dest.unlink(missing_ok=True) proc = subprocess.Popen( [exe, "-y", "-loglevel", "error", "-f", "rawvideo", "-pix_fmt", "gray12le", "-s", f"{width}x{height}", "-r", str(max(1.0, float(fps))), "-i", "pipe:0", "-c:v", "libx265", *depth_video_encoder_args(), "-movflags", "+faststart", str(tmp_dest)], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) written = 0 try: for frame_index in range(frame_count): codes = _frame_to_codes(frame_index) proc.stdin.write(codes.tobytes()) # little-endian uint16 written += 1 proc.stdin.close() stderr = proc.stderr.read() rc = proc.wait(timeout=1800) except Exception: proc.kill() tmp_dest.unlink(missing_ok=True) return 0, 0.0, 0, 0 if rc != 0 or not tmp_dest.exists() or tmp_dest.stat().st_size == 0: tmp_dest.unlink(missing_ok=True) print(f"[lerobot_export] depth video encode failed: {stderr.decode()[-300:]}") return 0, 0.0, 0, 0 tmp_dest.rename(dest) return written, max(1.0, float(fps)), height, width # 手套压力列 → 触觉语义命名:压力/触觉数据与视觉骨骼(hand_* / # observation.state.hand_*)完全分离。原始采集端文件仍为 # observation.left_glove/right_glove(协议不动),导出数据集统一为 # observation.tactile.left/right。 _TACTILE_COLUMN_MAP = { "observation.left_glove": "observation.tactile.left", "observation.right_glove": "observation.tactile.right", } def _read_sensor_rows(session_dir: Path, episode_id: str | None = None, skip_columns: set[str] | None = None) -> dict[int, dict]: """Merge independently stored glove/action parquet rows by frame_index. ``skip_columns``: 调用方已排除的列。这些列在读取阶段就被跳过 —— 触觉力矩阵每帧 0.8MB,整集 14MB,在 SSHFS 上读完要十几秒;既然最终 不会写进数据集,就不该读。仅按列名裁剪,不改变任何保留列的值。 """ rows: dict[int, dict] = {} try: import pyarrow.parquet as pq except ImportError as exc: raise RuntimeError("pyarrow is required for LeRobot export") from exc episode_index = _episode_storage_index(Path(session_dir), episode_id) for data_dir in session_dir.rglob("data"): if not data_dir.is_dir(): continue # The input root contains only the canonical project data tree, so # processing results cannot be mistaken for a second raw dataset. for path in sorted(data_dir.rglob("*.parquet")): if (episode_index is not None and path.name != f"episode_{episode_index:06d}.parquet"): continue try: if skip_columns: # 先读 schema(只读 footer,便宜),再只取需要的列。 available = [str(n) for n in pq.read_schema(path).names] wanted = [n for n in available if n not in skip_columns] table = pq.read_table(path, columns=wanted) else: table = pq.read_table(path) names = table.schema.names if "frame_index" not in names: continue for source_row in table.to_pylist(): fi = int(source_row.get("frame_index", 0)) target = rows.setdefault(fi, {}) # 传播帧过滤:对齐官方"只导出真实检测"语义(与 # _read_hand_3d_rows 的 state 过滤一致)。传播/插值 # 帧(检测短暂丢失时用前后帧插值延续显示)的手部 # 列置 None → _sanitize_row 按列语义补 NaN/False, # 训练端以 *_valid 列判定该帧该手无效。 sides, top_slots, devices = _propagated_hand_sets(source_row) for key, value in source_row.items(): if key == "frame_index" or key in {"timestamp", "episode_index"}: continue if key == "imu_ts_ns": # IMU 采样时间戳(ns,变长,每帧一批)→ 与 # observation.imu 同构透传,导出时走变长聚合 target["observation.imu_ts_ns"] = value continue if key.startswith("observation.") or key == "action" or key.startswith("action"): out_key = key if key.startswith("observation.") else "action" # 手套压力列重命名为触觉语义(与骨骼完全分离); # 保持 256 扁平(pa.Table.from_pylist 不支持 2 维值), # 16×16 语义在 features note 里声明 out_key = _TACTILE_COLUMN_MAP.get(out_key, out_key) if _is_propagated_hand_col(out_key, sides, top_slots, devices): target[out_key] = None continue target[out_key] = value except Exception: continue # The authoritative data directory is normally the first one. Keep # scanning sibling roots only when no frame data was found. if rows: break return rows # 旧 7 类手势索引 → 类名(2D mediapipe_hand 产物;3D 角度手势不经过此表) LEGACY_GESTURE_NAMES = { 0: "closed_fist", 1: "open_palm", 2: "pointing_up", 3: "thumbs_down", 4: "thumbs_up", 5: "victory", 6: "ilove_you", } # MediaPipe 21 关键点名称(OpenPose 手部顺序,与 TachinTactileGlove 一致) _HAND_KP_NAMES = ( "WRIST", "THUMB_CMC", "THUMB_MCP", "THUMB_IP", "THUMB_TIP", "INDEX_FINGER_MCP", "INDEX_FINGER_PIP", "INDEX_FINGER_DIP", "INDEX_FINGER_TIP", "MIDDLE_FINGER_MCP", "MIDDLE_FINGER_PIP", "MIDDLE_FINGER_DIP", "MIDDLE_FINGER_TIP", "RING_FINGER_MCP", "RING_FINGER_PIP", "RING_FINGER_DIP", "RING_FINGER_TIP", "PINKY_MCP", "PINKY_PIP", "PINKY_DIP", "PINKY_TIP", ) def _hand_xyz_names() -> list[str]: """63 维扁平手部坐标的逐元素 names:21 关键点 × xyz(官方逐元素 命名规范,参考 sensexperience pose/state 的 names 字段)。""" return [f"{kp}_{axis}" for kp in _HAND_KP_NAMES for axis in ("x", "y", "z")] def _tactile_names(size: int = 256) -> list[str]: """16×16 触觉阵列扁平化的逐元素 names(r<行>c<列>,行主序)。""" side = int(round(size ** 0.5)) if side * side != size: return [f"px_{i}" for i in range(size)] return [f"r{r}c{c}" for r in range(side) for c in range(side)] def _gesture_str(value) -> str: """手势统一为字符串标签:3D 角度判定标签原样;旧 7 类 int → 类名。""" if value is None or (isinstance(value, float) and value != value): # NaN return "" if isinstance(value, str): return value.strip() try: iv = int(value) except (TypeError, ValueError): return "" if iv < 0: return "" return LEGACY_GESTURE_NAMES.get(iv, str(iv)) def _hand_entry(source_row: dict, slot: str) -> dict | None: """从 hand_keypoints.parquet 行提取单手记录(无检测 → None)。""" import numpy as np if source_row.get(f"{slot}_2d_present") is False: return None kp = source_row.get(f"{slot}_keypoints") if kp is None: return None # 历史 hand_keypoints 产物是 21x3 嵌套列表,RGB_TO_3D 的 canonical # hand_3d 产物则是扁平 63 列。两种布局都属于同一套 2D 关键点合同。 try: raw = np.asarray(kp, dtype=np.float32) if raw.size != 63: return None arr = raw.reshape(21, 3) except (TypeError, ValueError): return None if arr.shape != (21, 3): return None # 不是 21×3,视为无效 if not np.isfinite(arr).all(): return None try: confidence = float(source_row.get(f"{slot}_confidence") or float("nan")) except (TypeError, ValueError): confidence = float("nan") return { "keypoints": arr.reshape(21, 3).tolist(), "handedness": str(source_row.get(f"{slot}_handedness") or source_row.get(f"{slot}_label") or ""), "confidence": confidence, # 2D 产物可能使用旧的 int 类别,也可能使用 RGB 估计模块的 # 空字符串/文本标签,统一成导出层的字符串合同。 "gesture": _gesture_str(source_row.get(f"{slot}_gesture")), "gesture_score": float(source_row.get(f"{slot}_gesture_score") or 0.0), } def _read_hand_rows(hand_paths: list[str], hand_3d_paths: list[str] | None = None) -> dict[int, dict]: """读手部骨骼产物 → {frame_index: {"hand_0": entry|None, "hand_1": entry|None}}。 ``hand_3d_paths`` 非空时走三角化分支(米制 3D 坐标,每个槽位直接覆盖 对应帧);否则走 2D 分支:每路读入,按槽位每帧取置信度最高的那路 (手部最清晰的画面,左右手可来自不同相机);entry 附带 ``source`` (相机 key,取自 parquet 文件名)供数据集标注来源列。单路行为与现状 一致,只是额外带上 source。 """ if hand_3d_paths: return _read_hand_3d_rows(hand_3d_paths) if not hand_paths: return {} import pandas as pd per_camera: list[dict[int, dict]] = [] for path in hand_paths: try: df = pd.read_parquet(path) except Exception: continue src = Path(path).stem # 相机 key(stereo_left / stereo_right) rows: dict[int, dict] = {} for _, source_row in df.iterrows(): fi = int(source_row.get("frame_index", 0)) rows[fi] = { "hand_0": _hand_entry(source_row, "hand_0"), "hand_1": _hand_entry(source_row, "hand_1"), } per_camera.append((src, rows)) merged: dict[int, dict] = {} for src, rows in per_camera: for fi, hand in rows.items(): target = merged.setdefault(fi, {"hand_0": None, "hand_1": None}) for slot in ("hand_0", "hand_1"): entry = hand.get(slot) if not entry: continue cur = target.get(slot) if cur is None or (entry.get("confidence") or 0) >= (cur.get("confidence") or 0): entry = dict(entry) entry["source"] = src target[slot] = entry return merged def _read_hand_3d_rows(hand_3d_paths: list[str]) -> dict[int, dict]: """读双目三角化产物(hand_3d/*.parquet)→ 与 2D 分支同构的 entry 映射。 每个槽位为三角化后的一只手(hand_0/1_label 记 handedness,landmarks_3d 为 63 = 21×3 米制 3D 坐标,左目相机系);entry 带 ``is_3d`` 标记供 features 声明区分 unit(meter vs normalized_image_coords)。 """ import numpy as np import pandas as pd # 函数内引用,模块顶层无全局 pd(缺失会 NameError) merged: dict[int, dict] = {} for path in hand_3d_paths: try: df = pd.read_parquet(path) except Exception: continue src = _hand_3d_source(path) # 物理 RGB source_key(旧产物回退文件名) for _, source_row in df.iterrows(): fi = int(source_row.get("frame_index", 0)) target = merged.setdefault(fi, {"hand_0": None, "hand_1": None}) for slot in ("hand_0", "hand_1"): if not source_row.get(f"{slot}_present"): continue # Black-glove D435 artifacts may contain bounded alpha-beta # predictions for preview continuity. Training exports are # real-only by default; the raw hand_3d parquet keeps the # propagated values and their explicit state fields. state = source_row.get(f"{slot}_state") if state not in (None, "", "real"): continue landmarks = source_row.get(f"{slot}_landmarks_3d") if landmarks is None: continue kp = np.asarray(landmarks, dtype=np.float64).reshape(21, 3) target[slot] = { "keypoints": kp.tolist(), "handedness": str(source_row.get(f"{slot}_label") or ""), # 新版 hand_3d 带 2D 检测置信度/角度手势列; # 旧产物缺列 → 1.0 / 空 / -1(语义保持兼容)。 # fingers 用显式 None 判断:0 = 握拳(合法值), # ``or -1`` 会把它错写成"无数据"再被垃圾列规则误删 "confidence": float(source_row.get(f"{slot}_confidence") or 1.0), "gesture": _gesture_str(source_row.get(f"{slot}_gesture")), "fingers": (int(source_row.get(f"{slot}_fingers")) if source_row.get(f"{slot}_fingers") is not None else -1), "gesture_score": 0.0, "source": src, "is_3d": True, "reprojection_error": float( source_row.get(f"{slot}_reprojection_error") or float("nan")), } return merged _RGB_ESTIMATED_3D_UNIT = "rgb_estimated_meters" _METRIC_3D_UNITS = {"camera_meters", "meter", "meters"} def _hand_3d_manifest(path: str | Path) -> dict: """读取 Hand3D 文件旁的 manifest,找不到时返回空字典。""" path = Path(path) candidates = [path.parent / f"{path.stem}.manifest.json"] candidates.extend(sorted(path.parent.glob(f"{path.stem}*.manifest.json"))) seen: set[Path] = set() for manifest in candidates: if manifest in seen: continue seen.add(manifest) try: value = json.loads(manifest.read_text(encoding="utf-8")) except (OSError, ValueError, TypeError): continue if isinstance(value, dict): return value return {} def _hand_3d_unit_for_path(path: str | Path) -> str | None: value = _hand_3d_manifest(path).get("unit") return str(value) if value else None def _is_rgb_estimated_hand_3d(path: str | Path) -> bool: """RGB_TO_3D 的 Hand3D 仅供预览,导出时必须降级为 2D。""" manifest = _hand_3d_manifest(path) unit = _hand_3d_unit_for_path(path) or "" mode = str(manifest.get("mode") or "").lower() return (unit == _RGB_ESTIMATED_3D_UNIT or "rgb_estimated_3d" in mode or "rgb_estimated" in mode) def _prepare_hand_export_inputs( hand_keypoints_paths: list[str] | None, hand_3d_paths: list[str] | None, hand_3d_right_paths: list[str] | None, hand_3d_unit: str | None = None) -> tuple[list[str], list[str], list[str], str | None]: """Separate preview-only RGB Hand3D from exportable hand artifacts. RGB_TO_3D writes a Hand3D parquet so the review UI can render a spatial skeleton. Its parquet also contains the original 2D landmarks, which are the only hand coordinates allowed into a dataset export. Only metric depth artifacts (``camera_meters``) stay on the 3D path. """ metric_paths: list[str] = [] estimated_paths: list[str] = [] for path in hand_3d_paths or []: (estimated_paths if _is_rgb_estimated_hand_3d(path) else metric_paths).append(str(path)) metric_right_paths: list[str] = [] estimated_right_paths: list[str] = [] for path in hand_3d_right_paths or []: (estimated_right_paths if _is_rgb_estimated_hand_3d(path) else metric_right_paths).append(str(path)) # RGB estimated Hand3D is read through its hand_*_keypoints columns. keypoint_paths = [str(path) for path in (hand_keypoints_paths or [])] keypoint_paths.extend(estimated_paths) keypoint_paths.extend(estimated_right_paths) # A unit supplied by an in-memory ArtifactRef may describe an RGB path # that was filtered out. Never let that stale value select a 3D schema. supplied_unit = str(hand_3d_unit or "") unit = supplied_unit if supplied_unit in _METRIC_3D_UNITS else None unit = unit or _infer_hand_3d_unit(metric_paths) return keypoint_paths, metric_paths, metric_right_paths, unit def _flat63(kps) -> list: """hand_*_world 统一为 63 位扁平列表(LeRobot 约定:多维列展平存储, features 声明 [21,3],加载器 reshape)。同时规避 pandas/pyarrow 对 嵌套列表的序列化怪癖(读回变 (21,) object 数组,下游 asarray 报错)。 """ import numpy as np if kps is None: return [float("nan")] * 63 arr = np.asarray(kps) if arr.dtype == object or (arr.ndim == 1 and arr.shape == (21,)): arr = np.asarray([np.asarray(p, dtype=np.float64) for p in kps]) arr = arr.astype(np.float64).reshape(-1) if arr.shape[0] != 63: arr = np.full(63, np.nan) return arr.tolist() def _force_matrix_semantics(root: Path) -> dict[str, dict]: """从源 canonical 数据集读力矩阵的语义元数据。 力矩阵列存的不是力值本身,而是**行内差分后的定标整数**:采集端先按规格 把 (250,250,3) 的力值放大 ``scale`` 倍(实测 100)再四舍五入成 int16, 然后逐行做差分压体积。还原要两步 —— 先 ``cumsum`` 再 ``÷ scale``。 这两个字段只写在源 ``meta/info.json`` 的 features 里,**parquet 本身 不带**。导出若不透传,产物就只剩一条裸整数列表:不知道要 cumsum(拿到 的一列差分值毫无意义),也不知道要除以 100(即使还原了也差两个数量级)。 数据一个字节没丢,但已经没人能正确解释它了。 ``encoding``/``scale``/``units``/``names`` 原样搬运,不做解释 —— 消费 者按源契约解码即可。 """ import json try: info = json.loads((Path(root) / "meta" / "info.json").read_text(encoding="utf-8")) except (OSError, ValueError): return {} features = info.get("features") if not isinstance(features, dict): return {} carried_keys = ("encoding", "scale", "units", "names") out: dict[str, dict] = {} for name, spec in features.items(): if "force_matrix" not in str(name) or not isinstance(spec, dict): continue carried = {k: spec[k] for k in carried_keys if k in spec} if carried: out[str(name)] = carried return out def _passthrough_dtype(key: str, sensor_rows: dict[int, dict]) -> str: """透传列的声明 dtype 必须与真实值一致(官方加载器按声明 cast, 实测 gesture 字符串列被泛型 float32 声明 → ArrowInvalid 崩溃)。 扫描前 50 帧的真实值推断:字符串 > bool > int > float。 """ import numpy as np seen_str = seen_bool = seen_int = False for frame_rows in list(sensor_rows.values())[:50]: value = frame_rows.get(key) if value is None or (isinstance(value, (float, np.floating)) and np.isnan(value)): continue if isinstance(value, str): seen_str = True elif isinstance(value, (bool, np.bool_)): seen_bool = True elif isinstance(value, (int, np.integer)): seen_int = True elif isinstance(value, (list, tuple, np.ndarray)): kind = np.asarray(value).dtype.kind if kind in "US": seen_str = True elif kind == "b": seen_bool = True elif kind in "iu": seen_int = True if seen_str: return "string" if seen_bool: return "bool" if seen_int: return "int64" return "float32" def _write_video_dense_keyframes(source: Path, destination: Path) -> None: """把视频写入数据集,必要时重编码以加密关键帧。 采集端用 ffmpeg 默认设置编码(GOP=250),10 秒的集只有 2 个关键帧 (0s 和 8.3s)。lerobot 按时间戳取帧时先 seek 到关键帧再向前解码,当 查询点紧邻下一个关键帧时 seek 会落到**目标之后**,该帧永远取不到, 训练直接抛 FrameTimestampError(实测 7 个采样点里坏 2 个)。 GOP=30(每秒一个关键帧)可让 seek 稳定落在目标之前,实测 7/7 全通过; 代价是文件约大 2 倍(2.7MB -> 5.7MB/集)。 编码失败(缺 ffmpeg、编码器不可用)时退回原样复制:宁可保留稀疏关键帧 也不要让整个导出失败 —— 取不到的帧会让那条样本报错,而不是产出错数据。 """ import subprocess try: result = subprocess.run( ["ffmpeg", "-y", "-v", "error", "-i", str(source), "-c:v", "libx264", "-g", "30", "-pix_fmt", "yuv420p", "-crf", "20", "-preset", "veryfast", str(destination)], capture_output=True, timeout=300) if result.returncode == 0 and destination.is_file() and destination.stat().st_size: return except (OSError, subprocess.SubprocessError): pass shutil.copy2(source, destination) def _ensure_observation_state(rows: list[dict]) -> str | None: """确保每帧有 ``observation.state``;缺失时由 UMI 列派生。 ACT 的 VAE 编码器无条件读取 ``batch[observation.state]`` (modeling_act.py 里那句 ``device=batch[OBS_STATE].device`` 没做 robot_state_feature 判空),缺这列直接 KeyError。即便抛开这个上游 问题,单帧观测的 ACT 也需要 proprioception 才知道夹爪在哪。 派生口径 ``[Δx, Δy, Δz, gripper]``:位置是本帧相对**本集第一帧**的 偏移。SLAM 的世界原点是每次采集初始化时任意确定的,直接喂绝对坐标 等于把无关的世界位置当特征。夹爪取开合度 0–1。 已有 observation.state 时不动(尊重上游的显式声明)。返回派生列名或 None。 """ import numpy as np if any("observation.state" in row for row in rows): return None if not any("observation.slam_pose" in row for row in rows): return None # 按 episode_index 分组:相对偏移必须在**各自集内**计算,跨集求差 # 会把两集世界原点的差异当成运动。 by_episode: dict[int, list[dict]] = {} for row in rows: try: episode = int(row.get("episode_index", 0)) except (TypeError, ValueError): episode = 0 by_episode.setdefault(episode, []).append(row) derived = 0 for episode_rows in by_episode.values(): origin = None for row in episode_rows: pose = row.get("observation.slam_pose") try: pose_values = np.asarray(pose, dtype=float).reshape(-1) except (TypeError, ValueError): continue if pose_values.size < 3 or not np.isfinite(pose_values[:3]).all(): continue if origin is None: origin = pose_values[:3].copy() gripper = 0.0 raw = row.get("observation.gripper_state") try: gripper_values = np.asarray(raw, dtype=float).reshape(-1) if gripper_values.size: gripper = float(np.clip(gripper_values[0] / 100.0, 0.0, 1.0)) except (TypeError, ValueError): pass row["observation.state"] = [ float(value) for value in np.concatenate([ pose_values[:3] - origin, [gripper]]) ] derived += 1 if not derived: return None print(f"[lerobot_export] derived observation.state for {derived} frames") return "observation.state" def _ensure_slam_pose(rows: list[dict]) -> str | None: """确保每帧有 ``observation.slam_pose``;缺失时由 ``slam_trajectory`` 重建。 采集端从 52 集起不再写 ``observation.slam_pose``,只写高频样本缓冲 ``observation.slam_trajectory``(旧批两列都有)。ACT 的输入特征里声明了 slam_pose,整列缺失会让官方加载器与策略直接 KeyError —— 这批数据因此 完全无法训练或评估。 重建口径见 ``app.umi_slam_action.poses_from_trajectory``(与派生 action 用的是同一个函数)。拿旧集(有 slam_pose)做过对照:位置相关系数 0.998~0.9999、逐帧位置差 std 仅 1~4mm(亚帧采样差),数值范围完全一致, 四元数 |dot| 中位 1.0000 —— **同一个物理量、同一个世界系**,不是近似。 重建版反而补上了采集端留的 28.5% 零填充行,质量更好。 起点偏移的口径与 ``_ensure_observation_state`` 一致,故两者可以串联:本 函数先跑,缺 state 的集还能继续从重建出的 slam_pose 派生 state。 已有该列时不动(旧批重导行为完全不变)。返回写入的列名或 None。 """ import numpy as np if any("observation.slam_pose" in row for row in rows): return None if not any("observation.slam_trajectory" in row for row in rows): return None from app.umi_slam_action import poses_from_trajectory # 按 episode_index 分组:时间基映射是每集单独拟合的(各集 SLAM 会话 # 独立),跨集混在一起会把不同集的时间基搅乱,重建出的位姿全错。 by_episode: dict[int, list[dict]] = {} for row in rows: try: episode = int(row.get("episode_index", 0)) except (TypeError, ValueError): episode = 0 by_episode.setdefault(episode, []).append(row) written = 0 for episode_rows in by_episode.values(): count = len(episode_rows) trajectories = [row.get("observation.slam_trajectory") for row in episode_rows] # 时间戳必须是稠密 float 数组:缺帧补 NaN(poses_from_trajectory 会 # 把它们标为缺失再插值),混进 None 会让 np.asarray 推断出 object。 stamps = np.full(count, np.nan, dtype=float) for index, row in enumerate(episode_rows): try: stamps[index] = float(row.get("timestamp")) except (TypeError, ValueError): continue poses = poses_from_trajectory(trajectories, stamps, count) if poses is None: continue for row, pose in zip(episode_rows, poses): row["observation.slam_pose"] = [float(value) for value in pose] written += 1 if not written: return None print(f"[lerobot_export] derived observation.slam_pose for {written} frames") return "observation.slam_pose" def _pad_sparse_matrix_columns(rows: list[dict], nan_fill: float) -> dict[str, int]: """定长矩阵/trajectory 列的空帧补齐到声明长度。 采集端在尚无数据的帧上写空数组(如触觉矩阵首帧、SLAM 丢帧时的 trajectory),同一列因此长短不一。官方加载器按声明的 shape 转成定长 Sequence,遇到空值直接 cast 失败: TypeError: Couldn't cast array of type list to List(Value('int64'), length=187500) 只补空值,不动有值帧 —— 缺失语义由列本身的值域表达。 返回 ``{列名: 补齐后的长度}``:features 声明必须用这个长度,而不是 首帧样本的长度 —— 首帧常常正是空的那一帧,拿它声明会得到 shape [0], 与实际写入的长度不符,加载器照样 cast 失败。 """ import numpy as np targets: dict[str, int] = {} for key in {k for row in rows for k in row}: if not (key.endswith("_matrix") or key.endswith("_trajectory")): continue target = 0 sample_value = None for row in rows: value = row.get(key) if isinstance(value, (list, tuple)) and len(value) > target: target = len(value) sample_value = value if not target or not sample_value: continue targets[key] = target # 补齐值的类型必须与该列已有元素**逐位一致**:触觉矩阵是 int16 # (声明 int64),补浮点会让 pyarrow 把整列推断成 double,与声明的 # dtype 冲突 → 加载器 cast 失败。int 列补 0,浮点列补 nan_fill。 existing = next((v for v in sample_value if v is not None), None) if isinstance(existing, (int, np.integer)) and not isinstance(existing, bool): fill = 0 else: fill = nan_fill for row in rows: value = row.get(key) if not isinstance(value, (list, tuple)) or len(value) == target: continue row[key] = list(value) + [fill] * (target - len(value)) return targets def _pad_variable_columns(rows: list[dict], nan_fill: float = float("nan")) -> tuple[int, dict[str, int]]: """变长/稀疏列固定长度化。返回 ``(imu_cap, {列名: 补齐长度})``。 - observation.imu:每样本 6 维 → 展平为 cap*6(nan_fill 补尾) - observation.imu_ts_ns:补 0 到 cap - 矩阵/trajectory 类的空帧:补齐到该列最大长度 (见 _pad_sparse_matrix_columns) 第二个返回值供 features 声明使用 —— 声明必须等于**实际写入的长度**, 否则官方加载器 cast 失败。 """ import numpy as np matrix_targets = _pad_sparse_matrix_columns(rows, nan_fill) imu_cap = 3 for row in rows: for key in ("observation.imu", "observation.imu_ts_ns"): value = row.get(key) if isinstance(value, (list, tuple)): imu_cap = max(imu_cap, len(value)) if not any("observation.imu" in row for row in rows): return imu_cap, matrix_targets for row in rows: imu = row.get("observation.imu") if isinstance(imu, (list, tuple)): samples: list[np.ndarray] = [] for sample in imu: if isinstance(sample, (list, tuple)): arr = np.asarray( [nan_fill if v is None else float(v) for v in sample], dtype=float) if arr.shape[0] < 6: arr = np.pad(arr, (0, 6 - arr.shape[0]), constant_values=nan_fill) samples.append(arr[:6]) else: samples.append(np.full(6, nan_fill)) while len(samples) < imu_cap: samples.append(np.full(6, nan_fill)) row["observation.imu"] = np.concatenate( samples, dtype=float).astype(float).tolist() ts = row.get("observation.imu_ts_ns") if isinstance(ts, (list, tuple)): row["observation.imu_ts_ns"] = ( [int(v) for v in ts] + [0] * (imu_cap - len(ts)))[:imu_cap] return imu_cap, matrix_targets def _propagated_hand_sets(source_row: dict) -> tuple[set[str], set[str], dict[str, set[str]]]: """该帧处于传播/插值状态(非 real)的手部槽位。 返回 (sides, top_slots, devices): - sides:{"left","right"} 中该帧为传播态的左右手(主命名空间 observation.state.hand__* 列据此过滤) - top_slots:{"0","1"} 顶层 observation.hand__* 列的传播槽位 - devices:{设备 token → 传播槽位集合},键取自 processing.hand_3d..hand__state,与 observation.state.devices..hand__* 同 token 旧会话无 state 列时全部视为 real(不过滤)。 """ sides: set[str] = set() top_slots: set[str] = set() devices: dict[str, set[str]] = {} for slot in ("0", "1"): state = source_row.get(f"hand_{slot}_state") if state not in (None, "", "real"): top_slots.add(slot) label = str(source_row.get(f"hand_{slot}_label") or "").lower() if label in ("left", "right"): sides.add(label) for key, value in source_row.items(): match = re.fullmatch(r"processing\.hand_3d\.(.+)\.hand_([01])_state", key) if match and value not in (None, "", "real"): devices.setdefault(match.group(1), set()).add(match.group(2)) return sides, top_slots, devices def _is_propagated_hand_col(key: str, sides: set[str], top_slots: set[str], devices: dict[str, set[str]]) -> bool: """手部列是否属于传播帧槽位(命中则整列置 None,消毒层补语义默认值)。""" match = re.fullmatch(r"observation\.hand_([01])_.*", key) if match: return match.group(1) in top_slots match = re.fullmatch(r"observation\.state\.devices\.([\w.]+)\.hand_([01])_.*", key) if match: return match.group(2) in devices.get(match.group(1), set()) match = re.fullmatch(r"observation\.state\.hand_(left|right)_.*", key) if match: return match.group(1) in sides return False def _sanitize_row(row: dict, nan_fill: float = float("nan")) -> None: """null 单元格兜底消毒(官方加载器对 null 零容忍:采样时 ``torch.tensor(None)`` 直接崩,实测上游 hand_3d 处理列带 null)。 缺失语义按列名填充:bool 列 False、字符串列空串、fingers -1、 其余(数值/列表)用 ``nan_fill``(v3.0 = NaN;v2.1 = 0.0,因为官方 v2.1→v3 转换脚本经 pandas/pyarrow 回写时会把 NaN 变 null,官方 加载器对 null 直接崩——缺失语义由 *_valid 列承载,不受影响); 嵌套列表内 None 同样置 nan_fill。 """ for key, value in list(row.items()): if isinstance(value, list): # ``None in value`` 走 C 层容器查找,比 ``any(v is None ...)`` # 的 Python 生成器快约 2.3 倍。语义等价:``in`` 用 ``==``, # 而数值/字符串与 None 比较恒为 False。 # 必须保留 isinstance 前置判断 —— numpy 数组的 ``__eq__`` # 返回数组,``in`` 会抛 "truth value of an array is ambiguous"。 if None in value: row[key] = [nan_fill if v is None else v for v in value] elif value is not None: continue elif key.endswith("_valid") or key == "next.done": row[key] = False elif (key.endswith(("_source", "_gesture", "_handedness")) or key in {"annotation", "annotation_scope"}): row[key] = "" elif key.endswith("_fingers"): row[key] = -1 else: row[key] = nan_fill def _junk_columns(rows: list[dict]) -> set[str]: """恒定哨兵垃圾列检测:数据从未真实写入的列不导出。 - reprojection_error:全程 NaN(手部管线未产生重投影误差) - fingers:全程 -1(手套传感器未接) - gesture_score:全程 0(角度判定路径无分数) - gesture:全程空串(无手势标签) 只要任一帧有真实值(非哨兵),整列保留。 """ import numpy as np if not rows: return set() junk: set[str] = set() for key in rows[0].keys(): vals = [row.get(key) for row in rows] if key.endswith("_reprojection_error"): arr = np.array([v for v in vals if v is not None], dtype=float) if arr.size == 0 or np.isnan(arr).all(): junk.add(key) elif key.endswith("_fingers"): if all(v == -1 for v in vals): junk.add(key) elif key.endswith("_gesture_score"): if all(v in (0.0, 0, None) for v in vals): junk.add(key) elif key.endswith("_gesture"): if all(not v for v in vals): junk.add(key) elif key.endswith("_handedness") or ( key.endswith("_source") and "_world" in key): # 恒定手部元数据:槽位左右手标记/坐标来源。单相机、手不换 # 位的批次全程恒定 → 与按 hand_left/right 拆分的列完全冗余; # 只要任一帧值不同(换手/多源混合)整列保留。 distinct = {v for v in vals if v} if len(distinct) <= 1: junk.add(key) elif key == "action": # 恒定 action = 全程同一动作值(采集端占位或机器人静止): # 对 imitation learning 零信息,与无 action 列的官方参考 # 行为一致;任一帧不同 → 真实动作数据,整列保留。 normalized = [tuple(v) if isinstance(v, (list, tuple)) else v for v in vals] if len(set(normalized)) <= 1: junk.add(key) elif isinstance(vals[0], (list, tuple)) and not any( isinstance(x, (list, tuple)) for x in vals[0]): # 定长标量向量列全程恒定(如手套 hand_pose 采集端占位全零) # = 零信息;真实手部/触觉数据在动/有接触,不可能全程恒定。 # imu 等嵌套列表(每帧样本批)跳过此规则。 normalized = [tuple(v) if isinstance(v, (list, tuple)) else v for v in vals] if len(set(normalized)) <= 1: junk.add(key) return junk def _cast_float32(rows: list[dict]) -> None: """浮点列统一 cast 成 float32。 pyarrow 对 Python float 推断 double;features 声明按官方参考用 float32(pose/timestamp 都是 float32),不 cast 会造成声明与实际 schema 不一致(lerobot-doctor Feature Consistency 的隐患)。 使用 np.float32 标量(而非 Python float),pyarrow 才会保留 float32。 """ import numpy as np for row in rows: for key, value in list(row.items()): if value is None or isinstance(value, (str, bool)): continue if isinstance(value, (list, tuple)): if value and all(isinstance(x, (list, tuple)) for x in value): # 嵌套列表(imu 每帧样本批):内层 6 维 cast float32 row[key] = [ [np.float32(v) if isinstance(v, (float, np.floating)) else v for v in sample] for sample in value ] elif value and all(isinstance(x, (float, np.floating)) and not isinstance(x, bool) for x in value): # 纯浮点列表(触觉/手部坐标等)cast float32; # int 列表(imu_ts_ns)保持 int64 row[key] = [np.float32(x) for x in value] elif isinstance(value, (float, np.floating)) and not isinstance(value, bool): row[key] = np.float32(value) def _hand_columns(hand: dict, suffix: str = "", coordinate_key: str = "world", device_namespace: str | None = None) -> dict: """手部骨骼 → LeRobot observation 列(缺检测 NaN 填充,shape 固定)。 左右世界坐标按 handedness 分配(hand_0 槽位未必是左手);handedness 未知时按槽位兜底(hand_0 → left, hand_1 → right)。 ``coordinate_key`` separates calibrated camera coordinates (``world``) from MediaPipe relative 3D (``3d``) and image-normalized 2D (``2d``). The old reader remains compatible with historical files; only newly exported datasets use the semantically correct field name. ``suffix``: 列名后缀,右目(辅助视角)数据用 "_rcam",与主列并存 (主数据 = 左目规范视图)。 ``device_namespace``: 多设备数据写入独立的 LeRobot/HDF5 命名空间: ``observation.state.devices..*``。 """ import numpy as np coordinate_key = coordinate_key if coordinate_key in {"world", "3d", "2d"} else "world" safe_namespace = (re.sub(r"[^A-Za-z0-9_.-]+", "_", str(device_namespace)).strip("._") if device_namespace else "") root = (f"observation.state.devices.{safe_namespace}" if safe_namespace else "") scalar_root = root or "observation" state_root = root or "observation.state" cols: dict = {} for slot in ("0", "1"): entry = hand.get(f"hand_{slot}") cols[f"{scalar_root}.hand_{slot}_gesture{suffix}"] = _gesture_str(entry["gesture"]) if entry else "" cols[f"{scalar_root}.hand_{slot}_fingers{suffix}"] = int(entry.get("fingers")) \ if entry and entry.get("fingers") is not None else -1 cols[f"{scalar_root}.hand_{slot}_gesture_score{suffix}"] = float(entry["gesture_score"]) if entry else 0.0 cols[f"{scalar_root}.hand_{slot}_confidence{suffix}"] = float(entry["confidence"]) if entry else float("nan") cols[f"{scalar_root}.hand_{slot}_handedness{suffix}"] = entry["handedness"] if entry else "" left = right = None left_src = right_src = None left_err = right_err = float("nan") for slot in ("hand_0", "hand_1"): entry = hand.get(slot) if not entry or entry["keypoints"] is None: continue h = entry["handedness"].lower() err = float(entry.get("reprojection_error") if entry.get("reprojection_error") is not None else float("nan")) if h == "left": left, left_src = entry["keypoints"], entry.get("source") left_err = err elif h == "right": right, right_src = entry["keypoints"], entry.get("source") right_err = err elif left is None: left, left_src = entry["keypoints"], entry.get("source") left_err = err elif right is None: right, right_src = entry["keypoints"], entry.get("source") right_err = err prefix = f"{state_root}.hand_left_{coordinate_key}{suffix}" cols[prefix] = _flat63(left) prefix = f"{state_root}.hand_right_{coordinate_key}{suffix}" cols[prefix] = _flat63(right) cols[f"{state_root}.hand_left_{coordinate_key}{suffix}_valid"] = left is not None cols[f"{state_root}.hand_right_{coordinate_key}{suffix}_valid"] = right is not None # 来源相机标注(多路识别时该帧关键点来自哪个相机;单路为空串) cols[f"{state_root}.hand_left_{coordinate_key}{suffix}_source"] = str(left_src or "") cols[f"{state_root}.hand_right_{coordinate_key}{suffix}_source"] = str(right_src or "") # 3D 关键点重投影误差(像素):质量评估/训练时按误差过滤低质量帧 cols[f"{state_root}.hand_left_{coordinate_key}{suffix}_reprojection_error"] = left_err cols[f"{state_root}.hand_right_{coordinate_key}{suffix}_reprojection_error"] = right_err return cols def _hand_3d_source(path: str | Path) -> str: """Resolve the physical RGB source from a Hand 3D manifest.""" path = Path(path) for manifest in sorted(path.parent.glob(f"{path.stem}*.manifest.json")): try: doc = json.loads(manifest.read_text(encoding="utf-8")) except (OSError, ValueError, TypeError): continue source = doc.get("source_key") or doc.get("rgb_source") \ or doc.get("left_source") if source: return str(source) return path.stem def _read_hand_3d_rows_by_source(paths: list[str] | None) -> dict[str, dict[int, dict]]: """Read independent per-device Hand 3D parquets without overwriting them.""" result: dict[str, dict[int, dict]] = {} for raw in paths or []: source = _hand_3d_source(raw) result[source] = _read_hand_3d_rows([str(raw)]) return result def _device_hand_features(source: str, coordinate_key: str, unit: str, note: str) -> dict[str, dict]: """Feature declarations matching ``device_namespace`` row columns.""" safe_source = re.sub(r"[^A-Za-z0-9_.-]+", "_", str(source)).strip("._") or "device" root = f"observation.state.devices.{safe_source}" features = {} for slot in ("0", "1"): features.update({ f"{root}.hand_{slot}_gesture": {"dtype": "string", "shape": [1]}, f"{root}.hand_{slot}_fingers": {"dtype": "int64", "shape": [1]}, f"{root}.hand_{slot}_gesture_score": {"dtype": "float32", "shape": [1]}, f"{root}.hand_{slot}_confidence": {"dtype": "float32", "shape": [1]}, f"{root}.hand_{slot}_handedness": {"dtype": "string", "shape": [1]}, }) for side in ("left", "right"): key = f"{root}.hand_{side}_{coordinate_key}" features.update({ key: {"dtype": "float32", "shape": [63], "names": _hand_xyz_names(), "unit": unit, "note": note}, f"{key}_valid": {"dtype": "bool", "shape": [1]}, f"{key}_source": {"dtype": "string", "shape": [1]}, f"{key}_reprojection_error": {"dtype": "float32", "shape": [1]}, }) return features def _infer_hand_3d_unit(paths: list[str] | None) -> str | None: """Read the coordinate unit from a hand_3d manifest for re-export. Re-export endpoints may not carry the original ArtifactRef metadata, so relying only on the in-memory metadata silently mislabeled depth output as MediaPipe-relative 3D. """ for raw_path in paths or []: path = Path(raw_path) for manifest in sorted(path.parent.glob("*.manifest.json")): try: value = json.loads(manifest.read_text(encoding="utf-8")).get("unit") except Exception: continue if value: return str(value) return None def _depth_assets(session_dir: Path, frame_count: int, episode_id: str | None = None) -> list[dict]: depth_root = _depth_root(session_dir, episode_id) if not depth_root.is_dir(): return [] assets = [] for source_dir in sorted(depth_root.iterdir()): if not source_dir.is_dir(): continue files = sorted(p for p in source_dir.rglob("*.png") if p.stem.isdigit()) if not files: continue indices = sorted({int(p.stem) for p in files}) depth_height = depth_width = 0 try: import cv2 import numpy as np sample = cv2.imdecode( np.fromfile(str(files[0]), dtype=np.uint8), cv2.IMREAD_UNCHANGED, ) if sample is not None: depth_height, depth_width = sample.shape[:2] except Exception: pass calibration: dict = {} calibration_dir = session_dir / "calibration" namespaced_calibration = (session_dir / "meta" / "calibration" / str(episode_id or "")) if episode_id and namespaced_calibration.is_dir(): calibration_dir = namespaced_calibration if not calibration_dir.is_dir(): calibration_dir = session_dir / "meta" / "calibration" if calibration_dir.is_dir(): for cal_path in sorted(calibration_dir.glob("*.json")): try: value = json.loads(cal_path.read_text(encoding="utf-8")) except Exception: continue if isinstance(value, dict): calibration.update({ key: value[key] for key in ("depth_scale", "depth_min", "depth_max") if key in value }) assets.append({ "source_key": source_dir.name, "dtype": "uint16", "unit": "meter" if calibration.get("depth_scale") else "raw", "depth_scale": calibration.get("depth_scale"), "frame_count": frame_count, "available_frame_count": len(indices), "missing_frames": [i for i in range(frame_count) if i not in set(indices)], "height": depth_height, "width": depth_width, **calibration, }) return assets # 内置默认双目标定 —— 批次标定缺失或整份全零时的兜底。键名取自该文件的 # ``name`` 字段(head_stereo),与会话里真有一份 head_stereo.json 时算出的键一致。 _DEFAULT_CALIBRATION_PATH = (Path(__file__).resolve().parent / "processing" / "modules" / "calibration_default_head_stereo.json") def _calibration_is_blank(doc: Any) -> bool: """整份标定里找不出任何非零数值 —— 视为未生效的空白标定。 只递归看数值叶子(标定本质上就是嵌套的数值结构);字符串等非数值字段 不构成"有效标定"的证据。所以 ``{"depth_scale": 0.001}`` 不算空白, 而全零的 intrinsic/rotation/translation 算。 """ if isinstance(doc, bool): return True # bool 是 int 子类,必须先挡掉 if isinstance(doc, (int, float)): return doc == 0 if isinstance(doc, dict): return all(_calibration_is_blank(v) for v in doc.values()) if isinstance(doc, (list, tuple)): return all(_calibration_is_blank(v) for v in doc) return True def _default_calibration_docs() -> dict[str, Any]: """读内置默认标定;文件缺失/损坏时返回 {} (调用方自行回落到空)。""" try: doc = json.loads(_DEFAULT_CALIBRATION_PATH.read_text(encoding="utf-8")) except (OSError, ValueError): return {} if not isinstance(doc, dict): return {} return {str(doc.get("name") or _DEFAULT_CALIBRATION_PATH.stem): doc} def _calibration_docs(session_dir: str | Path, episode_id: str | None = None) -> dict[str, Any]: """会话标定 → ``{名称: 标定内容}``,供导出侧写进数据集 meta。 目录口径与其余读取点一致(media_groups / depth_hand_3d / lerobot_v21 / 本文件的深度标定段):先看 ``calibration/``;给了 ``episode_id`` 时优先 ``meta/calibration//``,再回退 ``meta/calibration/``。键名优先 取 JSON 里的 ``name`` 字段,否则用文件名 stem。 **批次标定缺失、或整份全零时回退内置默认标定** (``processing/modules/calibration_default_head_stereo.json``)。不这么做的话, 导出的 meta 会声明一个从未生效的标定,下游照它复现 3D 全是错的。 只有**每一份**都空白才回退 —— 只要有一份像样的就用会话里的。 """ session_dir = Path(session_dir) calib_dir = session_dir / "calibration" if episode_id: namespaced = session_dir / "meta" / "calibration" / str(episode_id) if namespaced.is_dir(): calib_dir = namespaced if not calib_dir.is_dir(): calib_dir = session_dir / "meta" / "calibration" docs: dict[str, Any] = {} if calib_dir.is_dir(): for path in sorted(calib_dir.glob("*.json")): try: value = json.loads(path.read_text(encoding="utf-8")) except (OSError, ValueError, TypeError): continue if not isinstance(value, dict): continue docs[str(value.get("name") or path.stem)] = value if docs and not all(_calibration_is_blank(v) for v in docs.values()): return docs return _default_calibration_docs() or docs def _device_metadata(session_dir: Path) -> list[dict]: """Return physical-device declarations for export-side pairing metadata.""" for name in ("metadata.json", "meta/info.json"): path = Path(session_dir) / name if not path.is_file(): continue try: doc = json.loads(path.read_text(encoding="utf-8")) except (OSError, ValueError, TypeError): continue devices = doc.get("devices") or [] if isinstance(devices, dict): devices = [dict(value or {}, key=key) for key, value in devices.items()] if isinstance(devices, list): return [value for value in devices if isinstance(value, dict)] return [] def _annotation_maps(episode_id: str) -> tuple[dict[int, tuple[str | None, int, list[str]]], list[dict], int]: frame_map: dict[int, tuple[str | None, int, list[str]]] = {} definitions = [] skipped_candidates = 0 for index, segment in enumerate(list_annotations(episode_id)): # AI 候选段(未确认)不进入数据集 —— 只有人工确认后才导出 if str(segment.get("status") or "confirmed") in { "candidate", "pending_retry"}: skipped_candidates += 1 continue start = int(segment.get("start_frame_index", 0)) end = int(segment.get("end_frame_index", start)) scope = segment.get("source_scope") or ["episode"] if isinstance(scope, str): scope = [scope] definition = { "id": segment.get("id"), "label": segment.get("label", ""), "start_frame": start, "end_frame": end, "source_scope": scope, "notes": segment.get("notes"), "keyframes": segment.get("keyframes") or [], } definitions.append(definition) for frame in range(start, end + 1): frame_map.setdefault(frame, (segment.get("label"), index, scope)) return frame_map, definitions, skipped_candidates def _embed_info_metadata(output_dir: Path, info_json: dict) -> None: """把序列化 info 写入 parquet 的 schema metadata(huggingface 键, 官方 pusht v3 同款:data/ 与 meta/episodes/ 下全部 parquet)。 官方内嵌内容 = {"info": {"features": ...}}(仅 features),且必须 是 HF datasets Features 可解析的形态(每项带 _type: Value/Sequence 等)——直接内嵌 info.json 的 lerobot features 会让官方加载器在 load_episodes 阶段 Features.from_dict 崩溃(实测 CastError)。部分 查看器/加载器读内嵌 info 而非 meta/info.json。单文件失败不阻塞 导出。 """ import pyarrow as pa import pyarrow.parquet as pq def hf_feature(feature: dict) -> dict: """lerobot feature 声明 → HF Features 序列化形态。 标量(shape [1])→ Value;一维向量 → Sequence(Value);二维 → 嵌套 Sequence。与 datasets.Features.from_dict 的解析规则对齐。 """ dtype = str(feature.get("dtype") or "float32") shape = tuple(int(dim) for dim in (feature.get("shape") or [1])) node = {"dtype": dtype, "_type": "Value"} for dim in reversed(shape[1:]): node = {"feature": node, "length": dim, "_type": "Sequence"} if len(shape) >= 1 and shape[0] != 1: node = {"feature": node, "length": shape[0], "_type": "Sequence"} return node def arrow_fallback(arrow_type) -> dict: """parquet 里有列但 features 未声明的兜底(官方规则:多列=加载失败)。 形态对齐官方 pusht episodes 内嵌写法:list → 一层 Sequence (官方 stats/* 嵌套 list 也只写一层;与字段不匹配时 datasets 会回退按 arrow 类型生成,不影响加载)。 """ if pa.types.is_boolean(arrow_type): return {"dtype": "bool", "_type": "Value"} if pa.types.is_string(arrow_type): return {"dtype": "string", "_type": "Value"} if pa.types.is_integer(arrow_type): return {"dtype": "int64", "_type": "Value"} if pa.types.is_floating(arrow_type): return {"dtype": "float64", "_type": "Value"} if pa.types.is_list(arrow_type): inner = {"dtype": "string", "_type": "Value"} return {"feature": inner, "_type": "Sequence"} return {"dtype": "float32", "_type": "Value"} features = info_json.get("features") or {} for pattern in ("data/**/*.parquet", "meta/episodes/**/*.parquet"): for path in output_dir.glob(pattern): try: table = pq.read_table(path) columns = set(table.column_names) embedded: dict[str, dict] = {} for key, feature in features.items(): if key not in columns: # video 等流不在 parquet 里,官方内嵌同样不含 continue embedded[key] = hf_feature(feature) for name in columns - features.keys(): embedded[name] = arrow_fallback( table.schema.field(name).type) payload = json.dumps({"info": {"features": embedded}}, ensure_ascii=False) table = table.replace_schema_metadata({"huggingface": payload}) pq.write_table(table, path) except Exception as exc: print(f"[lerobot_export] embed info metadata failed " f"for {path}: {exc}") def _canonicalize_video_features(features: dict[str, dict]) -> None: """Mirror legacy ``video_info`` into the official ``feature.info`` block. The reference v3 reader still detects the legacy depth marker, but its depth decoder reads quantization parameters from ``features[*].info``. Keeping only ``video_info`` therefore silently falls back to the official 0.01–10 m defaults instead of this project's 100–5000 mm log domain. Preserve ``video_info`` for the application while writing the canonical fields used by current LeRobot readers. """ for feature in features.values(): if not isinstance(feature, dict) or feature.get("dtype") != "video": continue legacy = feature.get("video_info") if not isinstance(legacy, dict): continue info = dict(feature.get("info") or {}) for key, value in legacy.items(): if key.startswith("video.") or key == "has_audio": info.setdefault(key, value) is_depth = bool( legacy.get("video.is_depth_map") or legacy.get("is_depth_map") ) if is_depth: # Official DepthEncoderConfig uses metres; the project codec uses # the mathematically equivalent log domain in millimetres. depth_info = { "is_depth_map": True, "video.is_depth_map": True, "video.codec": "hevc", "video.pix_fmt": "gray12le", "video.channels": 1, "video.depth_min": DEPTH_MIN_MM / 1000.0, "video.depth_max": DEPTH_MAX_MM / 1000.0, "video.shift": 0.0, "video.use_log": True, "depth_unit": "mm", } info.update(depth_info) # v2.1 keeps this application-side block. Keep it truthful too: # otherwise older readers can interpret the single-plane stream # as a 3-channel image before they inspect the canonical v3 info. feature["video_info"] = { **legacy, "video.channels": 1, "video.is_depth_map": True, "video.is_depth_visualization": False, "video.codec": "hevc", "video.pix_fmt": "gray12le", "video.depth_min": DEPTH_MIN_MM / 1000.0, "video.depth_max": DEPTH_MAX_MM / 1000.0, "video.shift": 0.0, "video.use_log": True, "depth_unit": "mm", } shape = list(feature.get("shape") or []) if len(shape) == 3 and shape[2] in (0, 3): shape[2] = 1 feature["shape"] = shape else: info.setdefault("is_depth_map", False) feature["info"] = info def build_lerobot_dataset(dataset_name: str, episode_ids: list[str], output_dir: Path, split_ratio: float = 0.9, include_video_keys: list[str] | None = None, hand_keypoints_paths: list[str] | None = None, hand_3d_paths: list[str] | None = None, hand_3d_right_paths: list[str] | None = None, version: str = "v3.0", hand_3d_unit: str | None = None, exclude_columns: list[str] | None = None, progress_callback=None) -> Path: """Build a LeRobot dataset (``version`` = "v2.1" | "v3.0"). ``exclude_columns``: 可选,精确列名列表,从数据集中整体剔除。 用于丢掉训练用不到但体积/计算代价极高的列(如 250×250×3 的触觉 力矩阵:每帧 0.8MB,且逐元素分位数统计是导出耗时的主要来源)。 被剔除的列不会出现在 parquet、features 声明和 stats 里。 Depth sources are exported as their preview video (``is_depth_map``) plus the per-frame ``observation.depth..valid`` marker; the original 16-bit PNG sequence stays in the session directory and is not copied into the dataset. ``include_video_keys``: 可选,只导出这些 source_key 对应的视频 (工作流中"输入卡片 → 导出节点"的连接决定导出范围)。None = 全量 导出(导出页面/API 行为不变);空列表 = 不导出任何视频。 ``hand_keypoints_paths``: 可选,手部骨骼识别产物(hand_keypoints.parquet) 路径(工作流中处理节点连到导出节点时传入)。写入 observation.state.hand_left/right_2d;缺检测帧 NaN 填充,保持固定 shape。 ``hand_3d_paths``: 可选,Hand Skeleton 模块输出的 hand_3d/*.parquet。 深度抬升产物写入 observation.state.hand_*_world。RGB_TO_3D 产物 仅用于前端空间预览,导出时自动读取其中的 2D 关键点并写入 observation.state.hand_*_2d,不会写入任何估计 3D 坐标。 ``hand_3d_right_paths``: 可选,右目(辅助视角)骨骼产物 (hand_3d_right/*.parquet)。真实深度产物传入时以 ``_rcam`` 后缀列组 写入;RGB_TO_3D 的右视图则降级为 2D 关键点,不写入估计 3D。 ``version``: 数据集布局版本。v2.1 = 每 episode 一个 parquet/mp4 + 官方 legacy JSONL episode 索引;v3.0 = 合并分片(默认)。 两版本结构互不兼容,与官方 convert_dataset_v21_to_v30 脚本一致。 """ try: import pyarrow as pa import pyarrow.parquet as pq except ImportError as exc: raise RuntimeError("pyarrow is required for LeRobot export") from exc episodes = [get_episode(str(episode_id)) for episode_id in episode_ids] episodes = [episode for episode in episodes if episode] if not episodes: raise RuntimeError("No valid episodes selected for export") output_dir.mkdir(parents=True, exist_ok=True) is_v2 = str(version or "v3.0").lower().startswith("v2") chunks_size = 1000 # 官方 v2.1 每 chunk 的 episode 数(episode_chunk = index // 1000) rows: list[dict] = [] episode_rows: list[dict] = [] # v2.1:每集一个 parquet,记录该集 rows 的 [start, end) 范围用于切片 episode_row_ranges: list[tuple[int, int, int, int, int]] = [] total_videos = 0 task_by_description: dict[str, int] = {} task_descriptions: dict[int, str] = {} episode_video_meta_records: list[dict[str, tuple[int, float]]] = [] # Device metadata is used only to derive the dataset robot_type. It is # intentionally not exported as a standalone meta/devices.json file. device_metadata_by_episode: dict[str, list[dict]] = {} # 力矩阵的 scale/encoding 等只存在于源 info.json,按 key 收集后写进 # 导出产物的 features 声明(见 _force_matrix_semantics)。源根 → 元数据 # 加一层缓存:同项目几十集共用一份 info.json,不必每集重读。 force_matrix_meta: dict[str, dict] = {} _force_meta_cache: dict[str, dict[str, dict]] = {} depth_source_keys: set[str] = set() video_features: dict[str, dict] = {} # 手部骨骼产物(工作流中处理节点连到导出节点时传入)—— 按 frame 对齐。 # RGB_TO_3D 的 hand_3d 文件只为前端空间预览而存在;导出层把它拆回 # hand_*_keypoints(2D)。只有真实深度 hand_3d 才保留在 3D 分支。 (hand_keypoints_paths, hand_3d_paths, hand_3d_right_paths, hand_3d_unit) = _prepare_hand_export_inputs( hand_keypoints_paths, hand_3d_paths, hand_3d_right_paths, hand_3d_unit) hand_rows = _read_hand_rows(hand_keypoints_paths, hand_3d_paths) # 新版 Hand 3D 是每个物理设备一份 parquet。保留上面的主字段兼容 # 读取,同时把每个设备写入 observation.state.devices.。 device_hand_rows = _read_hand_3d_rows_by_source(hand_3d_paths) # 右目(辅助视角)手部骨骼:单独列组(_rcam 后缀),与主数据并存; # 未连接/无产物 → None,不写任何 _rcam 列(旧工作流行为不变)。 hand_rows_rcam = (_read_hand_3d_rows(hand_3d_right_paths) if hand_3d_right_paths else None) # Re-export endpoints may not carry ArtifactRef metadata. Read the # manifest so depth output keeps its camera/world semantics on rebuild. hand_3d_unit = hand_3d_unit or _infer_hand_3d_unit(hand_3d_paths) coordinate_key = ( "world" if hand_3d_unit in {"camera_meters", "meter", "meters"} else "3d" if hand_3d_paths else "2d" ) for out_episode_index, episode in enumerate(episodes): episode_id = str(episode["id"]) session_dir = Path(episode["path"]) device_metadata_by_episode[episode_id] = _device_metadata(session_dir) # 力矩阵元数据随源数据集而来。跨项目导出可能混到不同录制规格的集 # (scale 10/100/1000 互不兼容),一个 features 声明只能写一个值 —— # 以先出现的为准并出声,不要静默按第一份解读整批数据。 root_key = str(session_dir) if root_key not in _force_meta_cache: _force_meta_cache[root_key] = _force_matrix_semantics(session_dir) for fm_key, fm_spec in _force_meta_cache[root_key].items(): previous = force_matrix_meta.get(fm_key) if previous is None: force_matrix_meta[fm_key] = fm_spec elif previous != fm_spec: print(f"[Export] 力矩阵规格不一致 {fm_key}: episode {episode_id} " f"为 {fm_spec},此前为 {previous};沿用前者") streams = episode.get("camera_streams") or {} # 本集实际导出的每路视频(帧数, fps) → episodes parquet 的 # videos//from_timestamp 映射(官方 v3 加载器按时间戳定位帧) episode_video_meta: dict[str, tuple[int, float]] = {} sensor_rows = _read_sensor_rows( session_dir, episode_id, skip_columns=set(str(name) for name in exclude_columns or [])) annotation_map, annotation_defs, _skipped_candidates = _annotation_maps(episode_id) description = str(episode.get("project") or "").strip() if description not in task_by_description: task_id = len(task_by_description) task_by_description[description] = task_id task_descriptions[task_id] = description task_id = task_by_description[description] # 切片标注文本也进任务注册表:每段一个 task(全局按文本去重), # 帧级 task_index 指向该帧所属切片任务 → task-conditioned 训练 # 可按子任务切片取指令;未标注帧回退 episode 级任务。 for seg in annotation_defs: label = (seg.get("label") or "").strip() if label and label not in task_by_description: seg_task_id = len(task_by_description) task_by_description[label] = seg_task_id task_descriptions[seg_task_id] = label fps = float(episode.get("fps") or 30) frame_count = int(episode.get("frame_count") or 0) for source_key, stream in sorted(streams.items()): # 工作流连接驱动:只导出连接到导出节点的视频源(双向子串匹配, # 与 find_videos 一致);未连接的视频默认不导出 if include_video_keys is not None and not _key_matches_any(source_key, include_video_keys): continue source = _source_path(stream.get("path")) if source is None or not source.exists(): continue if is_v2: # v2.1 官方布局:videos/{video_key}/chunk-{episode_chunk:03d}/ episode_chunk = out_episode_index // chunks_size destination = (output_dir / "videos" / f"observation.images.{source_key}" / f"chunk-{episode_chunk:03d}") output_video = destination / f"episode_{out_episode_index:06d}.mp4" else: destination = (output_dir / "videos" / "chunk-000" / f"observation.images.{source_key}") output_video = destination / f"file-{out_episode_index:03d}.mp4" destination.mkdir(parents=True, exist_ok=True) _write_video_dense_keyframes(source, output_video) total_videos += 1 probed_count, probed_fps, width, height = _probe_video(source) episode_video_meta[f"observation.images.{source_key}"] = ( probed_count, probed_fps or fps) frame_count = max(frame_count, probed_count) fps = probed_fps or fps video_features[f"observation.images.{source_key}"] = { "dtype": "video", "shape": [height or 0, width or 0, 3], "names": ["height", "width", "channel"], "video_info": { "video.fps": fps, "video.is_depth_map": False, "has_audio": False, "source_episode": episode_id, }, } # 深度双路导出: # 1) {source} = JET 伪彩 h264(浏览器/网页 viewer 可播, # 给人看;HEVC 12bit 浏览器解码不了) # 2) {source}_metric = 官方同款 metric 深度视频(gray12le / # HEVC 12bit 无损,is_depth_map: true, # 官方加载器按量化参数反量化回米/毫米) # 无 PNG 真值的老批次只有 JET → 主键标 is_depth_map: true 兜底。 depth_assets = _depth_assets(session_dir, frame_count, episode_id) depth_source_keys.update(asset["source_key"] for asset in depth_assets) png_depth_sources = { asset["source_key"] for asset in depth_assets if any((session_dir / "depth" / asset["source_key"]).rglob("*.png")) or any(_depth_root(session_dir, episode_id) .joinpath(asset["source_key"]).rglob("*.png")) } generated_depth: set[str] = set() for asset in depth_assets: depth_source = asset["source_key"] if depth_source not in png_depth_sources: continue metric_key = f"observation.images.{depth_source}_metric" if is_v2: episode_chunk = out_episode_index // chunks_size destination = (output_dir / "videos" / metric_key / f"chunk-{episode_chunk:03d}") output_video = destination / f"episode_{out_episode_index:06d}.mp4" else: destination = (output_dir / "videos" / "chunk-000" / metric_key) output_video = destination / f"file-{out_episode_index:03d}.mp4" destination.mkdir(parents=True, exist_ok=True) png_dir = _depth_root(session_dir, episode_id) / depth_source pngs = sorted(p for p in png_dir.rglob("*.png") if p.stem.isdigit()) d_count, d_fps, d_height, d_width = _write_depth_video( output_video, pngs, frame_count, fps, asset.get("depth_scale")) if d_count <= 0: # 编码失败(缺 ffmpeg/libx265 等):本集 JET 主键将在下方 # mp4 循环以 is_depth_map: true 兜底,绝不静默丢深度 continue generated_depth.add(depth_source) total_videos += 1 episode_video_meta[metric_key] = (d_count, d_fps) video_features[metric_key] = { "dtype": "video", "shape": [d_height or 0, d_width or 0, 1], "names": ["height", "width", "channel"], "video_info": { "video.fps": d_fps, "video.height": d_height, "video.width": d_width, "video.channels": 1, "video.codec": "hevc", "video.pix_fmt": "gray12le", "video.is_depth_map": True, "video.depth_min_mm": DEPTH_MIN_MM, "video.depth_max_mm": DEPTH_MAX_MM, "video.depth_qmax": DEPTH_QMAX, "video.depth_qp": DEPTH_QP, "video.depth_quantization": "log", "video.depth_encoding": DEPTH_VIDEO_ENCODING, "video.depth_scale": asset.get("depth_scale") or 0.001, "has_audio": False, "source_episode": episode_id, "depth_source": depth_source, }, } # 深度双路导出: # 1) {source} = JET 伪彩 h264(浏览器/网页 viewer 可播, # 给人看;HEVC 12bit 浏览器解码不了) # 2) {source}_metric = 官方同款 metric 深度视频(gray12le / # HEVC 12bit 无损,is_depth_map: true, # 官方加载器按量化参数反量化回米/毫米) # 无 PNG 真值的老批次只有 JET → 主键标 is_depth_map: true 兜底。 for depth_source, depth_path in _depth_video_paths(session_dir, episode_id): if is_v2: episode_chunk = out_episode_index // chunks_size destination = (output_dir / "videos" / f"observation.images.{depth_source}" / f"chunk-{episode_chunk:03d}") output_video = destination / f"episode_{out_episode_index:06d}.mp4" else: destination = (output_dir / "videos" / "chunk-000" / f"observation.images.{depth_source}") output_video = destination / f"file-{out_episode_index:03d}.mp4" destination.mkdir(parents=True, exist_ok=True) _copy_depth_video_canonical(depth_path, output_video) total_videos += 1 depth_probe = _probe_depth_video(depth_path) _depth_count = int(depth_probe.get("frames") or 0) depth_fps = float(depth_probe.get("fps") or fps) depth_width = int(depth_probe.get("width") or 0) depth_height = int(depth_probe.get("height") or 0) # The canonical stream is a single-plane gray12le. Declare the # channels from the actual pixel format instead of the RGB-loop # default of 3 so the intermediate feature is truthful before # _canonicalize_video_features runs downstream. depth_channels = ( 1 if depth_probe.get("pix_fmt") in {"gray12le", "gray16le", "gray10le"} else 3 ) episode_video_meta[f"observation.images.{depth_source}"] = ( _depth_count, depth_fps or fps) has_metric = depth_source in generated_depth video_features[f"observation.images.{depth_source}"] = { "dtype": "video", "shape": [depth_height or 0, depth_width or 0, depth_channels], "names": ["height", "width", "channel"], "video_info": { "video.fps": depth_fps or fps, "video.height": depth_height or 0, "video.width": depth_width or 0, "video.channels": depth_channels, # 有 metric 流时主键是纯可视化;无 metric(老批次/ # 编码失败)才把主键当深度流兜底 "video.is_depth_map": not has_metric, "video.is_depth_visualization": has_metric, "has_audio": False, "source_episode": episode_id, "depth_source": depth_source, **({"note": "JET 伪彩可视化(h264,浏览器可播);" f"metric 真值见 {depth_source}_metric"} if has_metric else {}), }, } if frame_count <= 0: frame_count = max(sensor_rows.keys(), default=-1) + 1 # 触觉数据只保留 observation.tactile.left/right 的 256 维原始 # 压力数组。触觉 VIRIDIS 预览视频不是 LeRobot 必需项,训练包不 # 再生成 observation.images.tactile_left/right,避免把显示层数据 # 混入最小训练数据集。 for frame_index in range(frame_count): anno, anno_index, scope = annotation_map.get(frame_index, (None, -1, ["episode"])) # 帧级 task_index 优先指向该帧所属切片任务(标注文本在任务 # 注册表里的 id);无标注帧回退 episode 级任务。 frame_task = task_by_description.get(anno) if anno else None row = { "episode_index": out_episode_index, "frame_index": frame_index, "timestamp": frame_index / fps, "index": len(rows), "task_index": frame_task if frame_task is not None else task_id, "annotation": anno or "", "annotation_index": anno_index, "annotation_scope": json.dumps(scope, ensure_ascii=False), # 不再写 language_persistent:官方 v3 该列是结构化消息行 # 列表(role/content/style…,annotation pipeline 产物), # 不是标注文本字符串;逐帧文本已由 annotation 列承载。 "next.done": frame_index == frame_count - 1, } row.update(sensor_rows.get(frame_index, {})) hand = hand_rows.get(frame_index) if hand and not device_hand_rows: # 主命名空间 observation.state.hand_* 仅作兜底:有设备命名 # 空间产物(devices..hand_*)时不再重复写同一份数据 row.update(_hand_columns(hand, coordinate_key=coordinate_key)) for device_source, rows_by_frame in device_hand_rows.items(): device_hand = rows_by_frame.get(frame_index) if device_hand: row.update(_hand_columns( device_hand, coordinate_key=coordinate_key, device_namespace=device_source)) if hand_rows_rcam: hand_rcam = hand_rows_rcam.get(frame_index) if hand_rcam: row.update(_hand_columns(hand_rcam, suffix="_rcam", coordinate_key=coordinate_key)) for asset in depth_assets: row[f"observation.depth.{asset['source_key']}.valid"] = ( frame_index not in set(asset.get("missing_frames", [])) ) _sanitize_row(row, nan_fill=0.0 if is_v2 else float("nan")) rows.append(row) episode_row: dict[str, Any] = { "episode_index": out_episode_index, "episode_id": episode_id, "length": frame_count, "task_index": task_id, "tasks": [task_descriptions[task_id]], "split": "train" if out_episode_index < max(1, int(len(episodes) * split_ratio)) else "val", # 官方 v3.0 episode 索引列(checker 必需):全局行区间 = # 该集在合并 data parquet 中的 [from, to) 偏移;单文件数据集 # 分片/文件位置恒 0(官方列名带前缀,照写) "dataset_from_index": len(rows) - frame_count, "dataset_to_index": len(rows), "data/chunk_index": 0, "data/file_index": 0, "meta/episodes/chunk_index": 0, "meta/episodes/file_index": 0, } # 官方 v3.0 每路视频的时间戳映射(单目/双目/多目通用,按实际 # 导出的 video key 生成):from_timestamp 为该集在视频文件内的 # 起始秒,to_timestamp 为结束秒。本导出每集一个视频文件 → 从 0 # 到探测时长;帧在 [from, to) 区间内按 fps 换算帧号。 for vkey, (vcount, vfps) in sorted(episode_video_meta.items()): duration = (vcount / vfps) if (vcount > 0 and vfps) else 0.0 episode_row[f"videos/{vkey}/chunk_index"] = 0 episode_row[f"videos/{vkey}/file_index"] = out_episode_index episode_row[f"videos/{vkey}/from_timestamp"] = 0.0 episode_row[f"videos/{vkey}/to_timestamp"] = duration episode_rows.append(episode_row) # v2.1:该集 rows 切片范围(index, task_id, length, start, end) episode_row_ranges.append((out_episode_index, task_id, frame_count, len(rows) - frame_count, len(rows))) episode_video_meta_records.append(episode_video_meta) if progress_callback is not None: try: progress_callback(done=out_episode_index + 1, total_episodes=len(episodes)) except Exception: pass # 进度回调失败不影响导出本身 # 调用方指定剔除的列(如触觉力矩阵):在补齐/统计/写盘之前删掉, # 这样 parquet、features 声明与 stats 三处天然一致,不需要各自过滤。 if exclude_columns: dropped = set(str(name) for name in exclude_columns) if dropped: for row in rows: for key in dropped: row.pop(key, None) print(f"[lerobot_export] excluded columns: {sorted(dropped)}") # SLAM 位姿列:采集端 52 集起只写 slam_trajectory,不再写 slam_pose。 # 必须赶在 _pad_variable_columns 之前做 —— 补齐会把冲空帧填成 0,而 # 0 在 slam_trajectory 里是一个"合法样本"(t=0 的位姿),重建时会被 # 当成真实帧混进去。此处补出的 slam_pose 又是下面 observation.state # 的派生依据,故顺序不能颠倒。 _ensure_slam_pose(rows) # ACT 的观测需要 proprioception:单帧图像不足以让模型知道夹爪当前 # 在哪。数据集若已有 observation.state 则原样保留;否则在 UMI 会话上 # 由 SLAM 位姿 + 夹爪开合派生 [Δx, Δy, Δz, gripper](位置为相对本集 # 起点的偏移 —— SLAM 世界原点是每集任意初始化的,绝对坐标是噪声)。 _ensure_observation_state(rows) # 变长异步传感器列固定长度化:官方加载器不支持变长列(实测 # list 对声明 int64[1] 直接 cast 失败)。取全数据集最大 # 样本数补齐,声明 shape 同步为固定长度,变长语义保留在 note。 imu_cap, padded_lengths = _pad_variable_columns( rows, nan_fill=0.0 if is_v2 else float("nan")) # 收敛:删恒定哨兵垃圾列(全程 NaN/-1/0/空串 → 数据从未写入)。 # 有任一帧真实值则整列保留;features 里对应声明一并剔除。 junk_cols = _junk_columns(rows) if junk_cols: for row in rows: for key in junk_cols: row.pop(key, None) print(f"[lerobot_export] dropped constant sentinel columns: " f"{sorted(junk_cols)}") # 浮点列统一 float32(官方参考 pose/timestamp 均为 float32; # pyarrow 对 Python float 默认推断 double,会导致声明与 schema 不符) _cast_float32(rows) # robot_type 从物理设备声明推导(如 d435 / s80m),不再写死 unknown。 robot_kinds: list[str] = [] for docs in device_metadata_by_episode.values(): for doc in docs or []: kind = str(doc.get("kind") or "").strip() if kind and kind not in robot_kinds: robot_kinds.append(kind) robot_type = "+".join(robot_kinds) or "unknown" # 刚写进 data/ 的表顺手留给 _write_stats:省掉它把整个 data 目录从磁盘 # 再读回来一遍(见 _write_stats 的 data_table 参数说明)。 data_table = None if is_v2: # v2.1: each episode has one data parquet. The official v2.1 # converter reads the legacy JSONL files below; do not create a # nested v3-style meta/episodes parquet tree in a v2.1 export. episode_index_rows = [] _v2_tables = [] for ep_index, (_, task_id, length, start, end) in enumerate( episode_row_ranges): episode_chunk = ep_index // 1000 data_dir = output_dir / "data" / f"chunk-{episode_chunk:03d}" data_dir.mkdir(parents=True, exist_ok=True) episode_table = pa.Table.from_pylist(rows[start:end]) pq.write_table(episode_table, data_dir / f"episode_{ep_index:06d}.parquet") _v2_tables.append(episode_table) record = { "episode_index": ep_index, "tasks": [task_descriptions[task_id]], "task_index": task_id, "length": length, "data/chunk_index": episode_chunk, "data/file_index": ep_index, "dataset_from_index": start, "dataset_to_index": end, } vmeta = (episode_video_meta_records[ep_index] if ep_index < len(episode_video_meta_records) else {}) for vkey, (vcount, vfps) in sorted(vmeta.items()): duration = (vcount / vfps) if (vcount > 0 and vfps) else 0.0 record[f"videos/{vkey}/chunk_index"] = episode_chunk record[f"videos/{vkey}/file_index"] = ep_index record[f"videos/{vkey}/from_timestamp"] = 0.0 record[f"videos/{vkey}/to_timestamp"] = duration episode_index_rows.append(record) # 官方 v2.1 legacy 索引:convert_dataset_v21_to_v30 转换脚本 # 硬性读取 meta/episodes.jsonl(缺失直接 FileNotFoundError)。 meta_dir = output_dir / "meta" meta_dir.mkdir(parents=True, exist_ok=True) (meta_dir / "episodes.jsonl").write_text( "\n".join(json.dumps(record, ensure_ascii=False, default=_json_default) for record in episode_index_rows) + "\n", encoding="utf-8") # 逐集表的 schema 可能因某一集整列为空而不同(类型推断随值走), # concat 失败就退回磁盘重读 —— 统计量宁可慢,不能缺。 try: data_table = (pa.concat_tables(_v2_tables, promote=True) if _v2_tables else None) except Exception: data_table = None _write_episodes_stats_v21(output_dir, episode_row_ranges, rows, episode_video_meta_records, video_features) else: data_dir = output_dir / "data" / "chunk-000" data_dir.mkdir(parents=True, exist_ok=True) data_table = pa.Table.from_pylist(rows) pq.write_table(data_table, data_dir / "file-000.parquet") # Official v3 stores the complete episode table in one parquet file # (the writer may shard this directory for very large datasets). A # per-episode file here is readable by our UI but is not the official # v3 layout and makes third-party loaders miss the index. episodes_path = (output_dir / "meta" / "episodes" / "chunk-000" / "file-000.parquet") episodes_path.parent.mkdir(parents=True, exist_ok=True) pq.write_table(pa.Table.from_pylist(episode_rows), episodes_path) features = { **video_features, "episode_index": {"dtype": "int64", "shape": [1]}, "frame_index": {"dtype": "int64", "shape": [1]}, "timestamp": {"dtype": "float32", "shape": [1]}, "index": {"dtype": "int64", "shape": [1]}, "task_index": {"dtype": "int64", "shape": [1]}, "annotation": {"dtype": "string", "shape": [1]}, "annotation_index": {"dtype": "int64", "shape": [1]}, "annotation_scope": {"dtype": "string", "shape": [1]}, "next.done": {"dtype": "bool", "shape": [1]}, } # 手部骨骼 feature 声明(连接驱动:产物传入才写入)。 # 字段名必须和坐标语义一致:world=深度相机米制,3d=相对 3D,2d=图像坐标。 is_3d = any( entry and entry.get("is_3d") for hand in hand_rows.values() for entry in hand.values() ) if hand_rows else False if hand_rows and not device_hand_rows: # 主命名空间 feature 声明仅作兜底(与行写入规则一致: # 有设备命名空间产物时主命名空间不写列、不声明)。 # unit 语义:hand_3d 产物 manifest 显式声明的 unit(如深度图 # 抬升的 camera_meters)优先;否则 world_landmarks 相对 3D / # 2D 归一化坐标按 is_3d 区分(stereo 路径行为不变)。 world_unit = hand_3d_unit or ( "mediapipe_world_relative" if coordinate_key == "3d" else "normalized_image_coords") world_note = ( "深度图抬升相机系米制 3D(彩色相机系,米)" if coordinate_key == "world" else "RGB-only 相机相对 3D 估计(无深度真值,近似米制,不可视为测量值)" if hand_3d_unit == "rgb_estimated_meters" else "MediaPipe world_landmarks 相对 3D(手腕相对、任意尺度)" if coordinate_key == "3d" else "MediaPipe 21 关键点归一化图像坐标") left_key = f"observation.state.hand_left_{coordinate_key}" right_key = f"observation.state.hand_right_{coordinate_key}" features.update({ left_key: { "dtype": "float32", "shape": [63], "names": _hand_xyz_names(), "unit": world_unit, "note": world_note, }, right_key: { "dtype": "float32", "shape": [63], "names": _hand_xyz_names(), "unit": world_unit, "note": world_note, }, f"{left_key}_valid": {"dtype": "bool", "shape": [1]}, f"{right_key}_valid": {"dtype": "bool", "shape": [1]}, f"{left_key}_source": { "dtype": "string", "shape": [1], "note": "该帧左手骨骼关键点来源相机 key(多路识别;空 = 单路)", }, f"{right_key}_source": { "dtype": "string", "shape": [1], "note": "该帧右手骨骼关键点来源相机 key(多路识别;空 = 单路)", }, f"{left_key}_reprojection_error": { "dtype": "float32", "shape": [1], "note": "该帧左手 3D 关键点重投影误差(像素),质量评估/训练过滤用", }, f"{right_key}_reprojection_error": { "dtype": "float32", "shape": [1], "note": "该帧右手 3D 关键点重投影误差(像素),质量评估/训练过滤用", }, "observation.hand_0_gesture": { "dtype": "string", "shape": [1], "note": ("手指伸展标签:3D 三角化路径为角度判定" "(fist / open:index,middle ...);2D 路径为 7 类名" "(closed_fist/open_palm/...);空 = 无手/未知"), }, "observation.hand_1_gesture": { "dtype": "string", "shape": [1], "note": ("手指伸展标签:3D 三角化路径为角度判定" "(fist / open:index,middle ...);2D 路径为 7 类名" "(closed_fist/open_palm/...);空 = 无手/未知"), }, "observation.hand_0_fingers": { "dtype": "int64", "shape": [1], "note": ("手指伸展 bitmask(thumb=1,index=2,middle=4,ring=8," "pinky=16;0=fist;-1=未知/旧数据)"), }, "observation.hand_1_fingers": { "dtype": "int64", "shape": [1], "note": ("手指伸展 bitmask(thumb=1,index=2,middle=4,ring=8," "pinky=16;0=fist;-1=未知/旧数据)"), }, "observation.hand_0_gesture_score": { "dtype": "float32", "shape": [1], "note": "旧 7 类手势模型分数;3D 角度手势路径恒 0", }, "observation.hand_1_gesture_score": { "dtype": "float32", "shape": [1], "note": "旧 7 类手势模型分数;3D 角度手势路径恒 0", }, "observation.hand_0_confidence": {"dtype": "float32", "shape": [1]}, "observation.hand_1_confidence": {"dtype": "float32", "shape": [1]}, "observation.hand_0_handedness": {"dtype": "string", "shape": [1]}, "observation.hand_1_handedness": {"dtype": "string", "shape": [1]}, }) # 右目(辅助视角)手部骨骼 feature 声明(连接驱动:hand_3d_right 产物 # 传入才写入)。_rcam 列组与主列独立,训练时可按需选用/忽略。 if hand_rows_rcam: rcam_left_key = f"observation.state.hand_left_{coordinate_key}_rcam" rcam_right_key = f"observation.state.hand_right_{coordinate_key}_rcam" features.update({ rcam_left_key: { "dtype": "float32", "shape": [63], "names": _hand_xyz_names(), "unit": "mediapipe_world_relative", "note": "右目(辅助视角)相机检测的左手世界坐标;与主列并存,列组独立", }, rcam_right_key: { "dtype": "float32", "shape": [63], "names": _hand_xyz_names(), "unit": "mediapipe_world_relative", "note": "右目(辅助视角)相机检测的右手世界坐标;与主列并存,列组独立", }, f"{rcam_left_key}_valid": {"dtype": "bool", "shape": [1]}, f"{rcam_right_key}_valid": {"dtype": "bool", "shape": [1]}, f"{rcam_left_key}_source": { "dtype": "string", "shape": [1], "note": "该帧左手骨骼关键点来源相机 key(右目数据恒为右目产物名)", }, f"{rcam_right_key}_source": { "dtype": "string", "shape": [1], "note": "该帧右手骨骼关键点来源相机 key(右目数据恒为右目产物名)", }, f"{rcam_left_key}_reprojection_error": { "dtype": "float32", "shape": [1], "note": "右目左手 3D 关键点重投影误差(像素),质量评估/训练过滤用", }, f"{rcam_right_key}_reprojection_error": { "dtype": "float32", "shape": [1], "note": "右目右手 3D 关键点重投影误差(像素),质量评估/训练过滤用", }, "observation.hand_0_gesture_rcam": { "dtype": "string", "shape": [1], "note": "右目槽位 0 手势标签(角度判定,与主列同规则)", }, "observation.hand_1_gesture_rcam": { "dtype": "string", "shape": [1], "note": "右目槽位 1 手势标签(角度判定,与主列同规则)", }, "observation.hand_0_fingers_rcam": { "dtype": "int64", "shape": [1], "note": "右目槽位 0 手指伸展 bitmask(同主列规则)", }, "observation.hand_1_fingers_rcam": { "dtype": "int64", "shape": [1], "note": "右目槽位 1 手指伸展 bitmask(同主列规则)", }, "observation.hand_0_gesture_score_rcam": { "dtype": "float32", "shape": [1], "note": "右目槽位 0 手势分数(角度路径恒 0)", }, "observation.hand_1_gesture_score_rcam": { "dtype": "float32", "shape": [1], "note": "右目槽位 1 手势分数(角度路径恒 0)", }, "observation.hand_0_confidence_rcam": {"dtype": "float32", "shape": [1]}, "observation.hand_1_confidence_rcam": {"dtype": "float32", "shape": [1]}, "observation.hand_0_handedness_rcam": {"dtype": "string", "shape": [1]}, "observation.hand_1_handedness_rcam": {"dtype": "string", "shape": [1]}, }) if device_hand_rows: world_unit = hand_3d_unit or ( "mediapipe_world_relative" if coordinate_key == "3d" else "normalized_image_coords") world_note = ( "深度图抬升相机系米制 3D(每个物理设备独立坐标系,米)" if coordinate_key == "world" else "RGB-only 相机相对 3D 估计(每个物理设备独立,无深度真值)" if hand_3d_unit == "rgb_estimated_meters" else "MediaPipe world_landmarks 相对 3D(设备独立)" if coordinate_key == "3d" else "MediaPipe 21 关键点归一化图像坐标(设备独立)") for device_source in device_hand_rows: features.update(_device_hand_features( device_source, coordinate_key, world_unit, world_note)) # 传感器/触觉列声明(采集端上传,透传进数据集)—— # 列在 parquet 但 features 缺失会让 LeRobot 加载器读不完整。 # 手套压力 = observation.tactile.left/right(16×16 压力阵列, # 与视觉骨骼 observation.state.hand_* 完全分离)。 if sensor_rows: sample = next(iter(sensor_rows.values())) for key, value in sample.items(): if key in features: continue import numpy as np arr = np.asarray(value) shape = list(arr.shape) if arr.ndim else [1] # 补齐过的列必须以实际写入的长度声明。首帧样本往往正是空的那 # 一帧(触觉矩阵首帧、SLAM 丢帧),照它声明会得到 shape [0] 或 # 偏小的长度,与实际写入的数据不符 → 官方加载器 cast 失败。 if key in padded_lengths: shape = [padded_lengths[key]] if key in ("observation.tactile.left", "observation.tactile.right"): features[key] = { "dtype": "float32", "shape": [256], "names": _tactile_names(256), "note": ("手套压力/触觉阵列(16×16 行主序扁平化,float32," "与视频帧率对齐,每帧取该帧窗口最后一个样本);" "burst 双读样本仅隔 ~2ms,不做变长聚合"), } elif key in ("observation.left_hand_pose", "observation.right_hand_pose"): features[key] = { "dtype": "float32", "shape": [63], "names": _hand_xyz_names(), "note": ("采集端手套姿态(21 关键点×xyz,每帧取该帧窗口最后" "一个样本)。与相机关键点 " "observation.state.devices.*.hand_* 并存:手套" "姿态关节精度高且遮挡鲁棒,相机关键点绝对 3D 最准" ",训练端自选"), } elif key == "observation.imu_ts_ns": # 固定长度化(每帧补齐到 imu_cap,不足补 0); # 有效样本数以 observation.imu 对应段的 NaN 判定 features[key] = { "dtype": "int64", "shape": [imu_cap], "note": ("IMU 采样时间戳(ns),固定补齐至每帧最多 " f"{imu_cap} 个样本,不足帧尾部补 0;与 " "observation.imu 逐样本对齐"), } elif key == "observation.imu": features[key] = { "dtype": "float32", "shape": [imu_cap * 6], "note": ("IMU 采样批次:每帧固定 " f"{imu_cap}×6 维展平(3 轴加速度+3 轴角速度)," "不足帧尾部样本补 NaN,时间戳见 " "observation.imu_ts_ns"), } else: features[key] = {"dtype": _passthrough_dtype(key, sensor_rows), "shape": shape} # 力矩阵的行内差分/定标倍率不是能从数据本身看出来的属性, # 必须跟着声明一起走,否则产物不可解读(见 _force_matrix_semantics)。 semantics = force_matrix_meta.get(key) if semantics: features[key].update(semantics) if junk_cols: features = {k: v for k, v in features.items() if k not in junk_cols} # 调用方剔除的列同样从声明里去掉。features 由 sensor_rows 样本推导, # 而剔除只作用于 rows,不过滤这里会留下"声明了但数据里没有"的列, # 官方加载器读不出来会报错。 if exclude_columns: excluded = set(str(name) for name in exclude_columns) features = {k: v for k, v in features.items() if k not in excluded} # 导出的行里有 observation.state 就必须声明,否则官方加载器读不完整。 # 声明放在剔除之后:调用方显式排除的列不该被这里复活。 # 补出来的 slam_pose 同样必须声明:features 是从**源 parquet** 采样的 # (sensor_rows),而这列由 _ensure_slam_pose 在内存里写进 rows,采样 # 看不到它 —— 不补声明就是"数据里有列、声明里没有",官方加载器读不完整。 if rows and "observation.slam_pose" in rows[0] and "observation.slam_pose" not in features: features["observation.slam_pose"] = { "dtype": "float32", "shape": [7], "names": ["x", "y", "z", "qx", "qy", "qz", "qw"], "note": ("世界系绝对位姿(米 + 四元数 xyzw)。采集端 52 集起不再" "写入该列,导出时由 observation.slam_trajectory 的 50Hz" "样本重建(与采集端原始列实测位置相关 0.998+,同一世界系)。"), } if rows and "observation.state" in rows[0] and "observation.state" not in features: features["observation.state"] = { "dtype": "float32", "shape": [4], "names": ["delta_x", "delta_y", "delta_z", "gripper"], "note": ("proprioception:本帧位置相对**本集第一帧**的偏移(米)" "+ 夹爪开合(0–1)。SLAM 世界原点是每集任意初始化的," "故用相对量而非绝对坐标。"), } # 深度有效帧标记列(行循环写入,不声明会导致官方加载器读不完整) for depth_key in sorted(depth_source_keys): key = f"observation.depth.{depth_key}.valid" features[key] = { "dtype": "bool", "shape": [1], "note": "该帧深度真值(采集端 PNG)是否存在;false 帧的深度视频码值无效", } train_episode_count = sum(1 for item in episode_rows if item["split"] == "train") split_spec = {"train": f"0:{train_episode_count}"} if train_episode_count < len(episode_rows): split_spec["val"] = f"{train_episode_count}:{len(episode_rows)}" _canonicalize_video_features(features) if is_v2: # v2.1: episode 文件路径模板(episode_chunk/episode_index 占位),并带 # total_videos/total_chunks(v3.0 中已移除,转换脚本会删)。 info_json: dict[str, Any] = { "codebase_version": "v2.1", "robot_type": robot_type, "total_episodes": len(episodes), "total_frames": len(rows), "total_tasks": len(task_descriptions), "chunks_size": chunks_size, "data_files_size_in_mb": 100, "video_files_size_in_mb": 100, "fps": float(episodes[0].get("fps") or 30), "splits": split_spec, "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", "video_path": "videos/{video_key}/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.mp4", # Keep video_info in v2.1 as well. Removing it loses # video.is_depth_map and makes an exported depth stream # indistinguishable from RGB for downstream readers. "features": features, "total_videos": total_videos, "total_chunks": 1, "extensions": { "episodes_file": "meta/episodes.jsonl", "episodes_stats_file": "meta/episodes_stats.jsonl", }, } else: info_json = { "codebase_version": "v3.0", "robot_type": robot_type, "total_episodes": len(episodes), "total_frames": len(rows), "total_tasks": len(task_descriptions), "chunks_size": chunks_size, "data_files_size_in_mb": 100, "video_files_size_in_mb": 100, "fps": float(episodes[0].get("fps") or 30), "splits": split_spec, "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", "video_path": "videos/chunk-{chunk_index:03d}/{video_key}/file-{file_index:03d}.mp4", "features": features, "extensions": { "episodes_file": "meta/episodes/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", }, } _write_json(output_dir / "meta" / "info.json", info_json) # 官方 parquet 内嵌元数据(huggingface 键,官方 pusht 同款):部分 # 查看器/加载器读 parquet schema metadata 里的内嵌 info 而非 # meta/info.json —— 缺失时数据集校验报 total_episodes/total_frames # 无效(实测 viewer 报错场景)。 # 先算 stats 再嵌 parquet 元数据。两者读写互不相干(stats 只用到内存里 # 刚写好的 data_table + video_features + videos/*.mp4;embed 只重写 # data/ 与 meta/episodes/ 的 schema metadata),换序不改变任何产物。 # 换序是为了内存:_embed_info_metadata 自己会把整份 data parquet 读进 # 内存(v3 是单文件,50 集约 650MB),若此时 data_table 还活着,峰值 # 就凭空多出一整个数据集的量。算完立刻放掉引用。 _write_stats(output_dir, video_features, data_table) data_table = None _embed_info_metadata(output_dir, info_json) tasks_path = output_dir / "meta" / "tasks.jsonl" tasks_path.parent.mkdir(parents=True, exist_ok=True) if is_v2: with tasks_path.open("w", encoding="utf-8") as handle: for task_id, description in sorted(task_descriptions.items()): handle.write(json.dumps({"task_index": task_id, "task": description}, ensure_ascii=False) + "\n") else: # v3 exports use the official Parquet task table and do not carry the # internal tasks.jsonl compatibility file. tasks_path.unlink(missing_ok=True) _write_tasks_parquet(output_dir, task_descriptions) # Annotation segments and device provenance are represented by the # per-frame parquet columns and info/episode metadata. Do not create # the non-standard standalone meta/annotations.jsonl or meta/devices.json # companions in either v2.1 or v3.0 exports. (output_dir / "meta" / "annotations.jsonl").unlink(missing_ok=True) (output_dir / "meta" / "devices.json").unlink(missing_ok=True) return output_dir