DataFlowing / processor /app /lerobot_export.py
stouchinglhh
Claude
sync processor from GitLab c697b6e: 版本 1.19.4 / 前端 1.12.4 —— 服务端 i18n(首屏不再闪英文)、推理模块并入模型页、前端抽出 shared/、缓存键改用前端号
bb59211
Raw History Blame Contribute Delete
141 kB
"""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.<source>.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/<source>/``. 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.<src>.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<element: int64>
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_<side>_* 列据此过滤)
- top_slots:{"0","1"} 顶层 observation.hand_<n>_* 列的传播槽位
- devices:{设备 token → 传播槽位集合},键取自
processing.hand_3d.<token>.hand_<n>_state,与
observation.state.devices.<token>.hand_<n>_* 同 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.<source_key>.*``。
"""
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/<episode_id>/``,再回退 ``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.<source>.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.<source_key>。
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/<key>/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.<source>.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> 对声明 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