admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-11 03:52:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(datasets): make recompute_stats read-only safe and support image/video stats

Browse Source 
Recompute stats without modifying the source dataset by symlinking the large
immutable files (data/, videos/, images/) and copying only meta/ as writable
files. This avoids duplicating the dataset and works on read-only sources
(e.g. a mounted HF repo that isn't yours). Symlinking individual files keeps
push_to_hub working.

Also implement the previously-unfinished image/video stats recomputation: when
skip_image_video=False, per-episode image/video stats are recomputed by sampling
and decoding frames, mirroring compute_episode_stats.
This commit is contained in:
CarolinePascal
2026-07-10 14:35:22 +02:00
parent e40b58a8df
commit bf4c9174a8
2 changed files with 124 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/datasets/dataset_tools.py
+86 -2
View File
@@ -52,8 +52,11 @@ from lerobot.utils.utils import flatten_dict
from .aggregate import aggregate_datasets
from .compute_stats import (
aggregate_stats,
auto_downsample_height_width,
compute_episode_stats,
compute_relative_action_stats,
get_feature_stats,
sample_indices,
)
from .dataset_metadata import LeRobotDatasetMetadata
from .image_writer import write_image
@@ -77,6 +80,7 @@ from .utils import (
update_chunk_file_indices,
)
from .video_utils import (
decode_video_frames,
encode_video_frames,
reencode_video,
)
@@ -1559,6 +1563,82 @@ def modify_tasks(
return dataset
def _load_episode_image_frames(
dataset: LeRobotDataset,
key: str,
ep_idx: int,
frame_offsets: list[int],
is_depth: bool,
) -> np.ndarray:
"""Load sampled frames of an image feature for one episode as a (N, C, H, W) array."""
ep = dataset.meta.episodes[ep_idx]
from_idx = ep["dataset_from_index"]
column = dataset.hf_dataset.with_format(None).select_columns(key)
frames = []
for offset in frame_offsets:
img = column[from_idx + offset][key]
if is_depth:
arr = np.array(img)
if arr.ndim == 2:
arr = arr[np.newaxis, ...]
else:
arr = np.transpose(np.array(img.convert("RGB"), dtype=np.uint8), (2, 0, 1))
frames.append(auto_downsample_height_width(arr))
return np.stack(frames)
def _load_episode_video_frames(
dataset: LeRobotDataset,
key: str,
ep_idx: int,
frame_offsets: list[int],
is_depth: bool,
) -> np.ndarray:
"""Load sampled frames of a video feature for one episode as a (N, C, H, W) array."""
ep = dataset.meta.episodes[ep_idx]
video_path = dataset.root / dataset.meta.get_video_file_path(ep_idx, key)
from_timestamp = ep[f"videos/{key}/from_timestamp"]
timestamps = [from_timestamp + offset / dataset.meta.fps for offset in frame_offsets]
frames = decode_video_frames(
video_path, timestamps, dataset.tolerance_s, return_uint8=not is_depth, is_depth=is_depth
)
return np.stack([auto_downsample_height_width(frame) for frame in frames.numpy()])
def _compute_visual_episode_stats(
dataset: LeRobotDataset,
ep_idx: int,
visual_keys: list[str],
) -> dict:
"""Compute per-episode statistics for image/video features by sampling frames.
Mirrors the image/video branch of :func:`compute_episode_stats`: per-channel stats
are computed on downsampled sampled frames, then RGB stats are rescaled to [0, 1]
(depth maps keep their native units).
"""
ep_length = dataset.meta.episodes[ep_idx]["length"]
frame_offsets = sample_indices(ep_length)
ep_stats = {}
for key in visual_keys:
is_depth = key in dataset.meta.depth_keys
if dataset.meta.features[key]["dtype"] == "video":
frames = _load_episode_video_frames(dataset, key, ep_idx, frame_offsets, is_depth)
else:
frames = _load_episode_image_frames(dataset, key, ep_idx, frame_offsets, is_depth)
stats = get_feature_stats(frames, axis=(0, 2, 3), keepdims=True)
normalization_factor = 1.0 if is_depth else 255.0
ep_stats[key] = {
k: v if k == "count" else np.squeeze(v / normalization_factor, axis=0)
for k, v in stats.items()
}
return ep_stats
def recompute_stats(
dataset: LeRobotDataset,
skip_image_video: bool = True,
@@ -1572,7 +1652,9 @@ def recompute_stats(
Args:
dataset: The LeRobotDataset to recompute stats for.
skip_image_video: If True (default), only recompute stats for numeric features
(action, state, etc.) and keep existing image/video stats unchanged.
(action, state, etc.) and keep existing image/video stats unchanged. If False,
image/video stats are also recomputed by sampling and decoding frames from each
episode (this reads the image/video files, unlike the numeric-only path).
relative_action: If True, compute action stats in relative space by
iterating all valid action chunks and subtracting the current state.
This matches the normalization distribution the model sees during
@@ -1626,8 +1708,8 @@ def recompute_stats(
raise ValueError(f"No parquet files found in {data_dir}")
all_episode_stats = []
# TODO: enable image and video stats re-computation
numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]]
visual_keys = [k for k, v in features_to_compute.items() if v["dtype"] in ["image", "video"]]
for parquet_path in tqdm(parquet_files, desc="Computing stats from data files"):
df = pd.read_parquet(parquet_path)
@@ -1644,6 +1726,8 @@ def recompute_stats(
episode_data[key] = np.array(values)
ep_stats = compute_episode_stats(episode_data, features_to_compute)
if visual_keys:
ep_stats.update(_compute_visual_episode_stats(dataset, int(ep_idx), visual_keys))
all_episode_stats.append(ep_stats)
if features_to_compute and not all_episode_stats:
src/lerobot/scripts/lerobot_edit_dataset.py
+38 -3
View File
@@ -167,7 +167,9 @@ Show dataset information without feature details:
--operation.type info \
--operation.show_features false
Recompute dataset statistics (saves to lerobot/pusht_recomputed_stats by default):
Recompute dataset statistics (saves to lerobot/pusht_recomputed_stats by default). The source
dataset is never modified: large files are symlinked and only meta/ is copied, so this also works
on read-only source datasets:
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type recompute_stats
@@ -178,6 +180,12 @@ Recompute stats and save to a specific new repo_id:
--new_repo_id lerobot/pusht_new_stats \
--operation.type recompute_stats
Recompute stats including image/video features (samples and decodes frames from each episode):
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type recompute_stats \
--operation.skip_image_video false
Recompute stats in-place (overwrites original dataset stats):
lerobot-edit-dataset \
--repo_id lerobot/pusht \
@@ -377,6 +385,30 @@ def _resolve_io_paths(
return output_repo_id, input_path, output_path
def _reference_copy_dataset(input_root: Path, output_root: Path) -> None:
"""Create a lightweight copy of a dataset that never modifies the source.
The directory tree is recreated with real directories, and every file is
symlinked to its source counterpart so no data is duplicated and the source is
only ever read. Files under ``meta/`` are instead copied as real, writable files
so that stats/info can be rewritten without touching the original. Symlinking
individual files (rather than whole directories) keeps ``push_to_hub`` working,
since ``Path.glob`` follows file symlinks but does not descend into symlinked
directories. This makes the operation safe on read-only source datasets.
"""
for src in input_root.rglob("*"):
rel = src.relative_to(input_root)
dst = output_root / rel
if src.is_dir():
dst.mkdir(parents=True, exist_ok=True)
elif rel.parts[0] == "meta":
dst.parent.mkdir(parents=True, exist_ok=True)
shutil.copyfile(src, dst) # copyfile ignores source perms, so dst is writable
else:
dst.parent.mkdir(parents=True, exist_ok=True)
dst.symlink_to(src.resolve())
def get_output_path(
repo_id: str,
new_repo_id: str | None,
@@ -674,14 +706,17 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None:
)
dataset = LeRobotDataset(cfg.repo_id, root=input_root)
else:
logging.info(f"Copying dataset from {input_root} to {output_root}")
logging.info(f"Referencing dataset from {input_root} into {output_root} (source is left untouched)")
if output_root.exists():
backup_path = output_root.with_name(output_root.name + "_old")
logging.warning(f"Output directory {output_root} already exists. Moving to {backup_path}")
if backup_path.exists():
shutil.rmtree(backup_path)
shutil.move(output_root, backup_path)
shutil.copytree(input_root, output_root)
# recompute_stats only reads data/ and rewrites meta/stats.json, so symlink the
# large immutable files and copy only meta/. This avoids duplicating the dataset
# and works even when the source dataset is read-only.
_reference_copy_dataset(input_root, output_root)
dataset = LeRobotDataset(output_repo_id, root=output_root)
logging.info(f"Recomputing stats for {cfg.repo_id}")