mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 13:09:40 +00:00
feat(dataset filtering): adding support for VLM based dataset filtering following lerobot annotation pipeline style
This commit is contained in:
@@ -239,6 +239,56 @@ Every module is on by default and can be toggled independently (set to
|
||||
| `--vqa.restrict_to_default_camera` | `false` | Ground VQA only on `--vlm.camera_key` (else every camera). |
|
||||
| `--executor.episode_parallelism` | `16` | Episodes processed concurrently within each phase. |
|
||||
|
||||
## Camera-view curation
|
||||
|
||||
`lerobot-curate-cameras` is a separate, lightweight command that uses the same
|
||||
VLM backend for a **dataset-filtering / curation** pass. It downloads only the
|
||||
**first episode**, then for each camera view asks the VLM to:
|
||||
|
||||
1. **flag** whether the view is blurry / unusable, and
|
||||
2. **label** the view with a canonical name from a closed vocabulary
|
||||
(`top`, `wrist`, `front`, `bottom`, `left`, `right`, plus two-word combos
|
||||
like `left_wrist`).
|
||||
|
||||
It runs in one of two modes:
|
||||
|
||||
- `--mode=report` (default) — write the labels + verdicts into `meta/`
|
||||
(`meta/camera_curation.json` and a `curation` block on each camera in
|
||||
`meta/info.json`). Nothing is moved; this is the cheap triage pass and works
|
||||
for any dataset.
|
||||
- `--mode=rename` — apply the labels by renaming each camera key to
|
||||
`observation.images.<label>`. For **video** datasets this is a
|
||||
**download-free, server-side Hub commit**: the `videos/<key>/` files are moved
|
||||
with the Hub's LFS copy/delete (no video is downloaded or re-encoded), and only
|
||||
the small `meta/` files are edited.
|
||||
|
||||
```bash
|
||||
# Cheap, mutation-free triage (writes meta/camera_curation.json):
|
||||
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=report
|
||||
|
||||
# Apply the labels by renaming camera keys on a new branch (keeps `main` intact):
|
||||
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=rename --branch=curated
|
||||
|
||||
# Run the VLM decision on a GPU via HF Jobs (same --job.* flags as above):
|
||||
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=rename --job.target=h200
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- The Hub rename is **in place** on the source repo — the Hub does not support
|
||||
cross-repo LFS copies. Use `--branch` to commit to a branch so `main` is
|
||||
preserved.
|
||||
- **Image** datasets store frames inside the data parquet, so their rename can't
|
||||
avoid touching the data; the rename falls back to a local rewrite (via
|
||||
[`rename_features`](./using_dataset_tools#rename-features)). Prefer `--mode=report`
|
||||
for image datasets.
|
||||
- Views judged unusable are only flagged by default (still renamed). Pass
|
||||
`--drop_unusable=true` (local path) to remove them.
|
||||
|
||||
Key options: `--mode`, `--branch`, `--n_frames`, `--view_vocabulary`,
|
||||
`--allow_combos`, `--on_collision`, `--drop_unusable`, and the shared
|
||||
`--vlm.*` / `--job.*` flags documented above.
|
||||
|
||||
## Contributing new modules
|
||||
|
||||
The pipeline is built to grow, and **contributions are very welcome** —
|
||||
|
||||
@@ -89,6 +89,28 @@ lerobot-edit-dataset \
|
||||
--operation.feature_names "['observation.images.top']"
|
||||
```
|
||||
|
||||
#### Rename Features
|
||||
|
||||
Rename feature keys — typically to canonicalize camera views (e.g.
|
||||
`observation.images.cam_0` → `observation.images.left_wrist`). A rename changes
|
||||
no pixel data, so it is a cheap key-remap: it rewrites `meta/` (info features,
|
||||
episode `videos/*` and `stats/*` columns, `stats.json`), moves the
|
||||
`videos/<key>/` directory, and — for image datasets — renames the embedded
|
||||
image column. Videos are **not** re-encoded.
|
||||
|
||||
```bash
|
||||
# Rename one or more camera keys
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht \
|
||||
--operation.type rename_features \
|
||||
--operation.name_mapping '{"observation.images.cam_0": "observation.images.left_wrist"}'
|
||||
```
|
||||
|
||||
If two targets collide (e.g. two cameras both labeled `top`), the operation
|
||||
raises by default; pass `--operation.on_collision suffix` to disambiguate
|
||||
deterministically (`top`, `top_2`, …). To label camera views automatically with
|
||||
a VLM, see the [Annotation Pipeline](./annotation_pipeline#camera-view-curation).
|
||||
|
||||
#### Convert to Video
|
||||
|
||||
Convert an image-based dataset to video format, creating a new LeRobotDataset where images are stored as videos. This is useful for reducing storage requirements and improving data loading performance. The new dataset will have the exact same structure as the original, but with images encoded as MP4 videos in the proper LeRobot format.
|
||||
|
||||
@@ -356,6 +356,7 @@ lerobot-imgtransform-viz="lerobot.scripts.lerobot_imgtransform_viz:main"
|
||||
lerobot-edit-dataset="lerobot.scripts.lerobot_edit_dataset:main"
|
||||
lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main"
|
||||
lerobot-annotate="lerobot.scripts.lerobot_annotate:main"
|
||||
lerobot-curate-cameras="lerobot.scripts.lerobot_curate_cameras:main"
|
||||
lerobot-rollout="lerobot.scripts.lerobot_rollout:main"
|
||||
|
||||
# ---------------- Tool Configurations ----------------
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""VLM camera-view curation for LeRobot datasets.
|
||||
|
||||
For each dataset, the first episode is inspected by a vision-language model to
|
||||
(1) judge whether each camera view is blurry/unusable and (2) assign a canonical
|
||||
view label (``top``/``wrist``/``front``/…). The labels can then be applied by
|
||||
renaming the camera keys — for video datasets via a download-free, server-side
|
||||
Hub commit. Exposed as the ``lerobot-curate-cameras`` CLI.
|
||||
"""
|
||||
|
||||
from .config import DEFAULT_VIEW_VOCABULARY, CameraCurationConfig
|
||||
from .curator import (
|
||||
CameraVerdict,
|
||||
build_name_mapping,
|
||||
build_report,
|
||||
curate_cameras,
|
||||
is_valid_view_label,
|
||||
rename_camera_keys_on_hub,
|
||||
write_report,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_VIEW_VOCABULARY",
|
||||
"CameraCurationConfig",
|
||||
"CameraVerdict",
|
||||
"build_name_mapping",
|
||||
"build_report",
|
||||
"curate_cameras",
|
||||
"is_valid_view_label",
|
||||
"rename_camera_keys_on_hub",
|
||||
"write_report",
|
||||
]
|
||||
@@ -0,0 +1,80 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Config for ``lerobot-curate-cameras`` (VLM camera-view curation)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
from lerobot.annotations.steerable_pipeline.config import AnnotationJobConfig, VlmConfig
|
||||
|
||||
# The closed vocabulary of canonical camera-view labels. Combos are formed by
|
||||
# joining two of these with ``_`` (e.g. ``left_wrist``).
|
||||
DEFAULT_VIEW_VOCABULARY: tuple[str, ...] = ("top", "wrist", "front", "bottom", "left", "right")
|
||||
|
||||
|
||||
@dataclass
|
||||
class CameraCurationConfig:
|
||||
"""Top-level config for ``lerobot-curate-cameras``.
|
||||
|
||||
The VLM decision only ever reads the first episode (a cheap partial
|
||||
download). ``mode="report"`` writes the labels + quality verdicts into
|
||||
``meta/`` and moves nothing (works for any dataset). ``mode="rename"``
|
||||
additionally renames the camera keys to ``observation.images.<label>`` —
|
||||
for video datasets this is a server-side, download-free Hub commit.
|
||||
"""
|
||||
|
||||
# Hub dataset id (downloaded when ``root`` is unset) — also the rename target.
|
||||
repo_id: str | None = None
|
||||
# Local dataset directory (skips the Hub download).
|
||||
root: Path | None = None
|
||||
|
||||
# "report": write mapping + verdicts into meta/, no file moves.
|
||||
# "rename": physically rename camera keys to observation.images.<label>.
|
||||
mode: str = "report"
|
||||
|
||||
# Commit target branch for the Hub rename; keeps ``main`` intact when set.
|
||||
# None commits to the default branch.
|
||||
branch: str | None = None
|
||||
|
||||
# Episode inspected by the VLM (first episode by default).
|
||||
episode_index: int = 0
|
||||
# Frames sampled from that episode per camera and shown to the VLM.
|
||||
n_frames: int = 4
|
||||
|
||||
# Closed label vocabulary and whether two-token combos (left_wrist) are allowed.
|
||||
view_vocabulary: tuple[str, ...] = DEFAULT_VIEW_VOCABULARY
|
||||
allow_combos: bool = True
|
||||
|
||||
# "error" raises on colliding target labels; "suffix" disambiguates (top -> top_2).
|
||||
on_collision: str = "error"
|
||||
# Remove cameras judged unusable (default: only flag them, still rename).
|
||||
drop_unusable: bool = False
|
||||
|
||||
# Where to write the machine-readable report (default <root>/meta/camera_curation.json).
|
||||
report_path: Path | None = None
|
||||
|
||||
vlm: VlmConfig = field(default_factory=VlmConfig)
|
||||
job: AnnotationJobConfig = field(default_factory=AnnotationJobConfig)
|
||||
|
||||
seed: int = 1729
|
||||
# Keyframe decode backend forwarded to ``decode_video_frames`` (None = default).
|
||||
video_backend: str | None = None
|
||||
|
||||
# Upload the result (rename mode). Kept off by default so runs are dry.
|
||||
push_to_hub: bool = False
|
||||
push_commit_message: str | None = None
|
||||
@@ -0,0 +1,349 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Camera-view curation: per-camera VLM quality + label judgments and the
|
||||
lightweight (download-free) Hub rename that applies the chosen labels.
|
||||
|
||||
The decision (:func:`curate_cameras`) is a pure function of a
|
||||
``{camera_key: [frames]}`` map and a VLM client, so it unit-tests with a stub
|
||||
VLM and no dataset. The orchestrating CLI (``lerobot-curate-cameras``) samples
|
||||
those frames from the dataset's first episode.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from lerobot.annotations.steerable_pipeline.frames import to_image_blocks
|
||||
from lerobot.datasets.dataset_tools import _remap_camera_key_in_meta, _resolve_rename_collisions
|
||||
from lerobot.datasets.io_utils import load_info, write_info
|
||||
from lerobot.utils.io_utils import write_json
|
||||
|
||||
from .config import CameraCurationConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_PROMPT_PATH = Path(__file__).parent / "prompts" / "camera_curation.txt"
|
||||
|
||||
# The canonical prefix every curated camera key gets.
|
||||
OBS_IMAGE_PREFIX = "observation.images."
|
||||
|
||||
|
||||
@dataclass
|
||||
class CameraVerdict:
|
||||
"""One camera's VLM verdict."""
|
||||
|
||||
camera_key: str
|
||||
usable: bool
|
||||
view_label: str | None
|
||||
blur_reason: str | None = None
|
||||
confidence: float | None = None
|
||||
# Populated by ``build_name_mapping`` once collisions are resolved.
|
||||
proposed_new_key: str | None = None
|
||||
|
||||
|
||||
def _load_prompt() -> str:
|
||||
return _PROMPT_PATH.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def is_valid_view_label(label: str, vocabulary: tuple[str, ...], allow_combos: bool) -> bool:
|
||||
"""True if ``label`` is a single vocab word, or (when allowed) an underscore
|
||||
combo of at most two distinct vocab words."""
|
||||
if not label:
|
||||
return False
|
||||
tokens = label.split("_")
|
||||
if not allow_combos:
|
||||
return len(tokens) == 1 and tokens[0] in vocabulary
|
||||
if not (1 <= len(tokens) <= 2):
|
||||
return False
|
||||
return all(tok in vocabulary for tok in tokens) and len(set(tokens)) == len(tokens)
|
||||
|
||||
|
||||
def _build_messages(frames: list[Any], cfg: CameraCurationConfig) -> list[dict[str, Any]]:
|
||||
if cfg.allow_combos:
|
||||
combo_rule = (
|
||||
"You may combine at most two of these words with an underscore when "
|
||||
"one word is not precise enough (e.g. \"left_wrist\"). "
|
||||
)
|
||||
else:
|
||||
combo_rule = "Use exactly one of these words (no combinations). "
|
||||
prompt = _load_prompt().format(
|
||||
vocabulary=", ".join(cfg.view_vocabulary),
|
||||
combo_rule=combo_rule,
|
||||
)
|
||||
content = [*to_image_blocks(frames), {"type": "text", "text": prompt}]
|
||||
return [{"role": "user", "content": content}]
|
||||
|
||||
|
||||
def _parse_verdict(camera_key: str, result: Any, cfg: CameraCurationConfig) -> CameraVerdict:
|
||||
"""Turn a parsed VLM JSON object into a :class:`CameraVerdict` (defensively)."""
|
||||
if not isinstance(result, dict):
|
||||
return CameraVerdict(camera_key=camera_key, usable=True, view_label=None, blur_reason=None)
|
||||
|
||||
usable = bool(result.get("usable", True))
|
||||
blur_reason = result.get("blur_reason")
|
||||
blur_reason = str(blur_reason) if blur_reason else None
|
||||
|
||||
raw_label = result.get("view_label")
|
||||
label = str(raw_label).strip().lower().replace(" ", "_") if raw_label else ""
|
||||
view_label = label if is_valid_view_label(label, cfg.view_vocabulary, cfg.allow_combos) else None
|
||||
if raw_label and view_label is None:
|
||||
logger.warning(
|
||||
"camera %s: VLM returned view_label=%r which is not in the vocabulary %s; leaving unlabeled",
|
||||
camera_key,
|
||||
raw_label,
|
||||
cfg.view_vocabulary,
|
||||
)
|
||||
|
||||
confidence = result.get("confidence")
|
||||
try:
|
||||
confidence = float(confidence) if confidence is not None else None
|
||||
except (TypeError, ValueError):
|
||||
confidence = None
|
||||
|
||||
return CameraVerdict(
|
||||
camera_key=camera_key,
|
||||
usable=usable,
|
||||
view_label=view_label,
|
||||
blur_reason=blur_reason,
|
||||
confidence=confidence,
|
||||
)
|
||||
|
||||
|
||||
def curate_cameras(
|
||||
frames_by_camera: dict[str, list[Any]],
|
||||
cfg: CameraCurationConfig,
|
||||
vlm: Any,
|
||||
) -> list[CameraVerdict]:
|
||||
"""Judge each camera's quality + view label from a few sampled frames.
|
||||
|
||||
``frames_by_camera`` maps a camera key to a list of decoded frames (torch
|
||||
tensors or PIL images). One batched ``generate_json`` call is issued across
|
||||
all cameras. Cameras with no frames are still reported (usable, unlabeled)
|
||||
so the caller sees the full camera set.
|
||||
"""
|
||||
ordered_keys = list(frames_by_camera)
|
||||
callable_keys = [k for k in ordered_keys if frames_by_camera[k]]
|
||||
|
||||
verdicts: dict[str, CameraVerdict] = {
|
||||
k: CameraVerdict(camera_key=k, usable=True, view_label=None) for k in ordered_keys
|
||||
}
|
||||
|
||||
if callable_keys:
|
||||
messages_batch = [_build_messages(frames_by_camera[k], cfg) for k in callable_keys]
|
||||
results = vlm.generate_json(messages_batch)
|
||||
for key, result in zip(callable_keys, results, strict=True):
|
||||
verdicts[key] = _parse_verdict(key, result, cfg)
|
||||
|
||||
return [verdicts[k] for k in ordered_keys]
|
||||
|
||||
|
||||
def build_name_mapping(
|
||||
verdicts: list[CameraVerdict],
|
||||
existing_features: dict[str, dict],
|
||||
cfg: CameraCurationConfig,
|
||||
) -> dict[str, str]:
|
||||
"""Compute ``{old_key: observation.images.<label>}`` for labeled cameras.
|
||||
|
||||
Cameras without a valid label (or already at their canonical name) are
|
||||
skipped. Collisions are resolved with ``cfg.on_collision`` and the resolved
|
||||
target is written back onto each verdict's ``proposed_new_key``.
|
||||
"""
|
||||
desired: dict[str, str] = {}
|
||||
for v in verdicts:
|
||||
if v.view_label is None:
|
||||
continue
|
||||
target = f"{OBS_IMAGE_PREFIX}{v.view_label}"
|
||||
if target != v.camera_key:
|
||||
desired[v.camera_key] = target
|
||||
|
||||
if not desired:
|
||||
return {}
|
||||
|
||||
resolved = _resolve_rename_collisions(desired, existing_features, cfg.on_collision)
|
||||
by_key = {v.camera_key: v for v in verdicts}
|
||||
for old, new in resolved.items():
|
||||
by_key[old].proposed_new_key = new
|
||||
return resolved
|
||||
|
||||
|
||||
def build_report(
|
||||
verdicts: list[CameraVerdict],
|
||||
mapping: dict[str, str],
|
||||
cfg: CameraCurationConfig,
|
||||
) -> dict[str, Any]:
|
||||
"""Assemble the machine-readable curation report."""
|
||||
return {
|
||||
"repo_id": cfg.repo_id,
|
||||
"episode_index": cfg.episode_index,
|
||||
"view_vocabulary": list(cfg.view_vocabulary),
|
||||
"cameras": {
|
||||
v.camera_key: {
|
||||
"view_label": v.view_label,
|
||||
"usable": v.usable,
|
||||
"blur_reason": v.blur_reason,
|
||||
"confidence": v.confidence,
|
||||
"proposed_new_key": mapping.get(v.camera_key),
|
||||
}
|
||||
for v in verdicts
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def write_report(
|
||||
root: Path,
|
||||
verdicts: list[CameraVerdict],
|
||||
mapping: dict[str, str],
|
||||
cfg: CameraCurationConfig,
|
||||
) -> Path:
|
||||
"""Write ``meta/camera_curation.json`` and stamp verdicts into ``info.json``.
|
||||
|
||||
Stamping goes into each camera's ``features[key]["info"]["curation"]`` so the
|
||||
verdict travels with the dataset. Returns the report path.
|
||||
"""
|
||||
report = build_report(verdicts, mapping, cfg)
|
||||
default_report_path = root / "meta" / "camera_curation.json"
|
||||
report_path = Path(cfg.report_path) if cfg.report_path is not None else default_report_path
|
||||
report_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
write_json(report, report_path)
|
||||
|
||||
info = load_info(root)
|
||||
changed = False
|
||||
for v in verdicts:
|
||||
feature = info.features.get(v.camera_key)
|
||||
if feature is None:
|
||||
continue
|
||||
feature.setdefault("info", {})
|
||||
if feature["info"] is None:
|
||||
feature["info"] = {}
|
||||
feature["info"]["curation"] = {
|
||||
"view_label": v.view_label,
|
||||
"usable": v.usable,
|
||||
"blur_reason": v.blur_reason,
|
||||
"confidence": v.confidence,
|
||||
}
|
||||
changed = True
|
||||
if changed:
|
||||
write_info(info, root)
|
||||
|
||||
return report_path
|
||||
|
||||
|
||||
def _swap_key_in_path(path: str, old_key: str, new_key: str) -> str:
|
||||
"""Rewrite the ``<old_key>`` path segment of a ``videos/<key>/...`` repo path."""
|
||||
prefix = f"videos/{old_key}/"
|
||||
return f"videos/{new_key}/{path[len(prefix):]}" if path.startswith(prefix) else path
|
||||
|
||||
|
||||
def rename_camera_keys_on_hub(
|
||||
repo_id: str,
|
||||
name_mapping: dict[str, str],
|
||||
local_root: Path,
|
||||
*,
|
||||
revision: str | None = None,
|
||||
branch: str | None = None,
|
||||
token: str | None = None,
|
||||
commit_message: str | None = None,
|
||||
) -> Any:
|
||||
"""Rename camera keys on the Hub without downloading video data.
|
||||
|
||||
Edits the small ``meta/`` files locally (under ``local_root``, which must be
|
||||
a writable dataset root whose ``meta/`` is already present), then commits, in
|
||||
one atomic ``create_commit``: ``CommitOperationCopy`` + ``CommitOperationDelete``
|
||||
to move each ``videos/<old>/*`` LFS file server-side, and ``CommitOperationAdd``
|
||||
for the edited meta files. Renames in place on ``repo_id`` (cross-repo copies
|
||||
are unsupported); pass ``branch`` to commit to a branch and keep ``main`` intact.
|
||||
|
||||
Only video keys can be moved this way — reject swaps/cycles and image keys
|
||||
(handled by the local ``rename_features`` path instead).
|
||||
"""
|
||||
from huggingface_hub import CommitOperationAdd, CommitOperationCopy, CommitOperationDelete, HfApi
|
||||
|
||||
# A swap/cycle (a target that is also a source) cannot be expressed in a
|
||||
# single base-revision commit; defer to the local rename path.
|
||||
swaps = set(name_mapping.values()) & set(name_mapping)
|
||||
if swaps:
|
||||
raise NotImplementedError(
|
||||
f"Hub rename cannot swap keys in one commit (offending: {sorted(swaps)}); "
|
||||
"use the local rename_features path for swaps/cycles."
|
||||
)
|
||||
|
||||
# Determine which OLD keys are video-stored (only those have a videos/ tree)
|
||||
# BEFORE remapping the metadata.
|
||||
info = load_info(local_root)
|
||||
video_old_keys = {
|
||||
old for old in name_mapping if info.features.get(old, {}).get("dtype") == "video"
|
||||
}
|
||||
image_old_keys = {
|
||||
old for old in name_mapping if info.features.get(old, {}).get("dtype") == "image"
|
||||
}
|
||||
if image_old_keys:
|
||||
raise NotImplementedError(
|
||||
f"Hub rename cannot move image data stored in the data parquet (keys: {sorted(image_old_keys)}); "
|
||||
"use --mode report (metadata mapping) or the local rename_features path for image datasets."
|
||||
)
|
||||
|
||||
# 1. Rewrite meta/ locally (info features, episodes columns, stats keys).
|
||||
_remap_camera_key_in_meta(local_root, name_mapping)
|
||||
|
||||
api = HfApi(token=token)
|
||||
operations: list[Any] = []
|
||||
|
||||
# 2. Add the (small) meta files we just edited.
|
||||
meta_dir = local_root / "meta"
|
||||
meta_files = [meta_dir / "info.json"]
|
||||
stats_file = meta_dir / "stats.json"
|
||||
if stats_file.exists():
|
||||
meta_files.append(stats_file)
|
||||
meta_files.extend(sorted((meta_dir / "episodes").glob("*/*.parquet")))
|
||||
for fpath in meta_files:
|
||||
rel = fpath.relative_to(local_root).as_posix()
|
||||
operations.append(CommitOperationAdd(path_in_repo=rel, path_or_fileobj=str(fpath)))
|
||||
|
||||
# 3. Move video LFS files server-side (copy + delete), no download.
|
||||
repo_files = api.list_repo_files(repo_id, repo_type="dataset", revision=revision)
|
||||
n_moved = 0
|
||||
for old in video_old_keys:
|
||||
new = name_mapping[old]
|
||||
prefix = f"videos/{old}/"
|
||||
for f in repo_files:
|
||||
if f.startswith(prefix):
|
||||
operations.append(
|
||||
CommitOperationCopy(src_path_in_repo=f, path_in_repo=_swap_key_in_path(f, old, new))
|
||||
)
|
||||
operations.append(CommitOperationDelete(path_in_repo=f))
|
||||
n_moved += 1
|
||||
logger.info(
|
||||
"hub rename: moving %d video file(s) server-side across %d camera(s)",
|
||||
n_moved,
|
||||
len(video_old_keys),
|
||||
)
|
||||
|
||||
commit_info = api.create_commit(
|
||||
repo_id=repo_id,
|
||||
repo_type="dataset",
|
||||
operations=operations,
|
||||
revision=branch or revision,
|
||||
commit_message=commit_message or "curate: rename camera views (lerobot-curate-cameras)",
|
||||
)
|
||||
return commit_info
|
||||
|
||||
|
||||
def as_report_dict(verdicts: list[CameraVerdict]) -> list[dict[str, Any]]:
|
||||
"""Convenience: verdicts as plain dicts (for logging/JSON)."""
|
||||
return [asdict(v) for v in verdicts]
|
||||
@@ -0,0 +1,28 @@
|
||||
You are inspecting frames from ONE camera of a robot manipulation dataset. All
|
||||
frames come from the same fixed camera during a single episode; use them
|
||||
together to judge the camera, not any single moment.
|
||||
|
||||
Do two things and return them as one JSON object.
|
||||
|
||||
1. QUALITY. Decide whether this camera view is usable for training a policy.
|
||||
Mark it UNUSABLE if it is blurry / out of focus, badly over- or
|
||||
under-exposed, mostly occluded, static/frozen, corrupted, or otherwise does
|
||||
not clearly show the scene. Otherwise it is usable.
|
||||
|
||||
2. VIEW LABEL. Choose the single best label for where this camera is mounted /
|
||||
what it looks at, using ONLY this closed vocabulary:
|
||||
|
||||
{vocabulary}
|
||||
|
||||
{combo_rule}Pick the label that best matches the viewpoint (e.g. a
|
||||
downward overhead shot is "top"; a camera on the robot's gripper/hand that
|
||||
moves with the arm is "wrist"). Do not invent words outside the vocabulary.
|
||||
|
||||
Output strictly valid JSON, no prose, no code fences, with exactly these keys:
|
||||
|
||||
{{
|
||||
"usable": true or false,
|
||||
"blur_reason": "<short reason if unusable, else null>",
|
||||
"view_label": "<one label from the vocabulary, combos joined by '_'>",
|
||||
"confidence": <number between 0 and 1>
|
||||
}}
|
||||
@@ -33,6 +33,7 @@ from .dataset_tools import (
|
||||
recompute_stats,
|
||||
reencode_dataset,
|
||||
remove_feature,
|
||||
rename_features,
|
||||
split_dataset,
|
||||
)
|
||||
from .factory import make_dataset, make_train_eval_datasets, resolve_delta_timestamps
|
||||
@@ -96,6 +97,7 @@ __all__ = [
|
||||
"recompute_stats",
|
||||
"reencode_dataset",
|
||||
"remove_feature",
|
||||
"rename_features",
|
||||
"resolve_delta_timestamps",
|
||||
"safe_stop_image_writer",
|
||||
"split_dataset",
|
||||
|
||||
@@ -47,6 +47,7 @@ from lerobot.configs import (
|
||||
)
|
||||
from lerobot.configs.video import DEPTH_ENCODER_INFO_FIELD_NAMES
|
||||
from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE
|
||||
from lerobot.utils.io_utils import load_json, write_json
|
||||
from lerobot.utils.utils import flatten_dict
|
||||
|
||||
from .aggregate import aggregate_datasets
|
||||
@@ -60,6 +61,8 @@ from .image_writer import write_image
|
||||
from .io_utils import (
|
||||
get_parquet_file_size_in_mb,
|
||||
load_episodes,
|
||||
load_info,
|
||||
to_parquet_one_row_group_per_episode,
|
||||
write_info,
|
||||
write_stats,
|
||||
write_tasks,
|
||||
@@ -72,7 +75,9 @@ from .utils import (
|
||||
DEFAULT_DATA_PATH,
|
||||
DEFAULT_EPISODES_PATH,
|
||||
DEPTH_FILE_PATTERN,
|
||||
EPISODES_DIR,
|
||||
IMAGE_FILE_PATTERN,
|
||||
STATS_PATH,
|
||||
VIDEO_DIR,
|
||||
update_chunk_file_indices,
|
||||
)
|
||||
@@ -484,6 +489,250 @@ def remove_feature(
|
||||
)
|
||||
|
||||
|
||||
# Columns in ``meta/episodes/*.parquet`` are namespaced by feature key under
|
||||
# these prefixes (e.g. ``videos/observation.images.top/from_timestamp`` and
|
||||
# ``stats/observation.images.top/mean``). Renaming a feature means rewriting the
|
||||
# middle ``<key>`` segment of every such column. Note ``stats/*`` columns are
|
||||
# invisible via ``meta.episodes`` (``load_episodes`` drops them), so we operate
|
||||
# on the raw parquet.
|
||||
_EPISODE_KEY_PREFIXES = ("videos", "stats")
|
||||
|
||||
# Features that must never be renamed (or become a rename target): the dataset
|
||||
# indexing/bookkeeping columns.
|
||||
_REQUIRED_FEATURES = frozenset({"timestamp", "frame_index", "episode_index", "index", "task_index"})
|
||||
|
||||
|
||||
def _resolve_rename_collisions(
|
||||
name_mapping: dict[str, str],
|
||||
existing_features: dict[str, dict],
|
||||
on_collision: str,
|
||||
) -> dict[str, str]:
|
||||
"""Validate/disambiguate a ``{old_key: new_key}`` mapping against collisions.
|
||||
|
||||
The post-rename key set is ``(features \\ sources) ∪ targets``. A collision is
|
||||
either two sources mapping to the same target, or a target equal to an
|
||||
untouched existing key. Swaps/cycles between sources are *not* collisions
|
||||
(handled downstream). ``on_collision="error"`` raises listing every offending
|
||||
pair; ``"suffix"`` disambiguates deterministically (``top`` → ``top_2`` → …)
|
||||
in sorted-source order.
|
||||
"""
|
||||
if on_collision not in ("error", "suffix"):
|
||||
raise ValueError(f"on_collision must be 'error' or 'suffix', got {on_collision!r}")
|
||||
|
||||
sources = set(name_mapping)
|
||||
untouched = set(existing_features) - sources
|
||||
targets = list(name_mapping.values())
|
||||
duplicate_targets = {t for t in targets if targets.count(t) > 1}
|
||||
untouched_collisions = set(targets) & untouched
|
||||
|
||||
if on_collision == "error":
|
||||
problems = []
|
||||
if duplicate_targets:
|
||||
problems.append(f"multiple cameras map to the same target(s): {sorted(duplicate_targets)}")
|
||||
if untouched_collisions:
|
||||
problems.append(
|
||||
f"target(s) collide with existing feature(s) not being renamed: "
|
||||
f"{sorted(untouched_collisions)}"
|
||||
)
|
||||
if problems:
|
||||
raise ValueError(
|
||||
"rename_features collision(s): "
|
||||
+ "; ".join(problems)
|
||||
+ ". Resolve the labels (e.g. use combos like 'left_wrist') or pass "
|
||||
"on_collision='suffix'."
|
||||
)
|
||||
return dict(name_mapping)
|
||||
|
||||
# suffix mode: greedily de-collide in a deterministic (sorted) order.
|
||||
used = set(untouched)
|
||||
resolved: dict[str, str] = {}
|
||||
for src in sorted(name_mapping):
|
||||
target = name_mapping[src]
|
||||
if target in used:
|
||||
base, i = target, 2
|
||||
while target in used:
|
||||
target = f"{base}_{i}"
|
||||
i += 1
|
||||
resolved[src] = target
|
||||
used.add(target)
|
||||
return resolved
|
||||
|
||||
|
||||
def _remap_camera_key_in_meta(root: Path, name_mapping: dict[str, str]) -> None:
|
||||
"""Rename feature keys across the dataset's ``meta/`` files (no file moves).
|
||||
|
||||
Touches: ``meta/info.json`` ``features`` (key renamed, feature dict carried
|
||||
verbatim so codec ``info`` / depth params survive), every
|
||||
``meta/episodes/*/*.parquet`` (``videos/<old>/*`` and ``stats/<old>/*``
|
||||
columns), and ``meta/stats.json`` (top-level ``<old>`` key). All three are
|
||||
simultaneous relabels, so swaps/cycles are safe here.
|
||||
"""
|
||||
# info.json — rebuild features preserving insertion order.
|
||||
info = load_info(root)
|
||||
info.features = {name_mapping.get(key, key): ft for key, ft in info.features.items()}
|
||||
write_info(info, root)
|
||||
|
||||
# episodes parquet — rename namespaced columns by prefix.
|
||||
def _rename_column(col: str) -> str:
|
||||
for prefix in _EPISODE_KEY_PREFIXES:
|
||||
head = f"{prefix}/"
|
||||
if col.startswith(head):
|
||||
rest = col[len(head) :]
|
||||
for old, new in name_mapping.items():
|
||||
if rest == old or rest.startswith(f"{old}/"):
|
||||
return f"{head}{new}{rest[len(old) :]}"
|
||||
return col
|
||||
|
||||
for path in sorted((root / EPISODES_DIR).glob("*/*.parquet")):
|
||||
df = pd.read_parquet(path)
|
||||
col_map = {c: _rename_column(c) for c in df.columns if _rename_column(c) != c}
|
||||
if col_map:
|
||||
df = df.rename(columns=col_map)
|
||||
to_parquet_one_row_group_per_episode(df, path)
|
||||
|
||||
# stats.json — remap top-level feature keys.
|
||||
stats_path = root / STATS_PATH
|
||||
if stats_path.exists():
|
||||
stats = load_json(stats_path)
|
||||
if isinstance(stats, dict):
|
||||
stats = {name_mapping.get(key, key): value for key, value in stats.items()}
|
||||
write_json(stats, stats_path)
|
||||
|
||||
|
||||
def _move_camera_key_dirs(root: Path, name_mapping: dict[str, str]) -> None:
|
||||
"""Move ``videos/<old>`` and ``images/<old>`` trees to their new key names.
|
||||
|
||||
Two-phase (source → sentinel → target) so a swap like ``{a: b, b: a}`` cannot
|
||||
clobber. Missing source dirs are skipped (a key may be stored one way only).
|
||||
"""
|
||||
for subdir in (VIDEO_DIR, "images"):
|
||||
base = root / subdir
|
||||
if not base.exists():
|
||||
continue
|
||||
# Phase 1: move every source to a unique sentinel.
|
||||
sentinels: dict[str, Path] = {}
|
||||
for i, old in enumerate(name_mapping):
|
||||
src = base / old
|
||||
if src.exists():
|
||||
sentinel = base / f".__rename_tmp_{i}__"
|
||||
shutil.move(str(src), str(sentinel))
|
||||
sentinels[old] = sentinel
|
||||
# Phase 2: sentinel → final target.
|
||||
for old, sentinel in sentinels.items():
|
||||
shutil.move(str(sentinel), str(base / name_mapping[old]))
|
||||
|
||||
|
||||
def _rename_image_data_columns(root: Path, name_mapping: dict[str, str]) -> None:
|
||||
"""Rename image-feature columns inside ``data/*.parquet`` at the Arrow level.
|
||||
|
||||
Image datasets embed frames as HF ``Image()`` columns in the data parquet.
|
||||
We rename the Arrow field *and* the matching key in the schema-level
|
||||
``huggingface`` metadata (which references columns by name), so no pixel
|
||||
bytes are decoded or re-embedded and ``datasets`` still types the column as
|
||||
an image after the rename.
|
||||
"""
|
||||
import json
|
||||
|
||||
data_dir = root / DATA_DIR
|
||||
if not data_dir.exists():
|
||||
return
|
||||
for path in sorted(data_dir.glob("*/*.parquet")):
|
||||
table = pq.read_table(path)
|
||||
col_map = {c: name_mapping[c] for c in table.column_names if c in name_mapping}
|
||||
if not col_map:
|
||||
continue
|
||||
table = table.rename_columns([col_map.get(c, c) for c in table.column_names])
|
||||
metadata = dict(table.schema.metadata or {})
|
||||
hf_key = b"huggingface"
|
||||
if hf_key in metadata:
|
||||
hf_meta = json.loads(metadata[hf_key])
|
||||
features = hf_meta.get("info", {}).get("features")
|
||||
if isinstance(features, dict):
|
||||
for old, new in col_map.items():
|
||||
if old in features:
|
||||
features[new] = features.pop(old)
|
||||
metadata[hf_key] = json.dumps(hf_meta).encode()
|
||||
table = table.replace_schema_metadata(metadata)
|
||||
pq.write_table(table, str(path))
|
||||
|
||||
|
||||
def rename_features(
|
||||
dataset: LeRobotDataset,
|
||||
name_mapping: dict[str, str],
|
||||
output_dir: str | Path | None = None,
|
||||
repo_id: str | None = None,
|
||||
*,
|
||||
on_collision: str = "error",
|
||||
) -> LeRobotDataset:
|
||||
"""Rename dataset feature keys without re-encoding any pixel data.
|
||||
|
||||
A rename changes zero frame content, so this does a cheap key-remap rather
|
||||
than the full-copy ``modify_features`` path (which would re-embed images and
|
||||
byte-copy videos). It rewrites ``meta/`` (info features, episodes
|
||||
``videos/*``+``stats/*`` columns, stats.json keys), moves the physical
|
||||
``videos/<key>/`` (and ``images/<key>/``) directories, and — for image
|
||||
datasets — renames the embedded ``data/*.parquet`` image column at the Arrow
|
||||
level. Feature ``info`` dicts (video codec params, depth ``is_depth_map``) are
|
||||
carried verbatim.
|
||||
|
||||
Args:
|
||||
dataset: The source LeRobotDataset.
|
||||
name_mapping: ``{old_feature_key: new_feature_key}``. Identity pairs are
|
||||
ignored. Typically used to canonicalize camera keys, e.g.
|
||||
``{"observation.images.cam_0": "observation.images.left_wrist"}``.
|
||||
output_dir: Where the renamed dataset is written. Defaults to
|
||||
``$HF_LEROBOT_HOME/repo_id``. When it equals ``dataset.root`` the
|
||||
rename is applied in place.
|
||||
repo_id: Identifier for the renamed dataset (default ``<repo_id>_renamed``).
|
||||
on_collision: ``"error"`` (default) raises on colliding targets;
|
||||
``"suffix"`` disambiguates deterministically (``top`` → ``top_2``).
|
||||
|
||||
Returns:
|
||||
The renamed LeRobotDataset.
|
||||
"""
|
||||
if not name_mapping:
|
||||
raise ValueError("name_mapping must be a non-empty {old_key: new_key} dict")
|
||||
|
||||
features = dataset.meta.features
|
||||
mapping = {old: new for old, new in name_mapping.items() if old != new}
|
||||
if not mapping:
|
||||
raise ValueError("name_mapping only contains identity renames (old == new); nothing to do")
|
||||
|
||||
missing = [old for old in mapping if old not in features]
|
||||
if missing:
|
||||
raise ValueError(f"Feature(s) not found in dataset: {missing}")
|
||||
|
||||
bad_required = sorted(
|
||||
{name for pair in mapping.items() for name in pair if name in _REQUIRED_FEATURES}
|
||||
)
|
||||
if bad_required:
|
||||
raise ValueError(f"Cannot rename to/from required features: {bad_required}")
|
||||
|
||||
bad_names = [new for new in mapping.values() if "/" in new]
|
||||
if bad_names:
|
||||
raise ValueError(f"Target feature name(s) cannot contain '/': {bad_names}")
|
||||
|
||||
mapping = _resolve_rename_collisions(mapping, features, on_collision)
|
||||
|
||||
if repo_id is None:
|
||||
repo_id = f"{dataset.repo_id}_renamed"
|
||||
output_dir = Path(output_dir) if output_dir is not None else HF_LEROBOT_HOME / repo_id
|
||||
|
||||
in_place = output_dir.resolve() == Path(dataset.root).resolve()
|
||||
if not in_place:
|
||||
shutil.copytree(dataset.root, output_dir)
|
||||
|
||||
image_keys = set(dataset.meta.image_keys)
|
||||
|
||||
_remap_camera_key_in_meta(output_dir, mapping)
|
||||
_move_camera_key_dirs(output_dir, mapping)
|
||||
image_mapping = {old: new for old, new in mapping.items() if old in image_keys}
|
||||
if image_mapping:
|
||||
_rename_image_data_columns(output_dir, image_mapping)
|
||||
|
||||
return LeRobotDataset(repo_id=repo_id, root=output_dir)
|
||||
|
||||
|
||||
def _fractions_to_episode_indices(
|
||||
total_episodes: int,
|
||||
splits: dict[str, float],
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Run ``lerobot-curate-cameras`` on HF Jobs (HuggingFace GPUs).
|
||||
|
||||
Same shape as the annotation submitter (``lerobot.jobs.annotate``): the VLM
|
||||
decision needs a GPU, so the pod boots the ``vllm/vllm-openai`` image, installs
|
||||
lerobot on top, and replays the user's CLI with ``lerobot-curate-cameras``. The
|
||||
``--mode=rename`` commit runs from the pod (which holds ``HF_TOKEN``); a bare
|
||||
``--mode=report`` run leaves its output only on the pod, so we warn.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shlex
|
||||
import sys
|
||||
from dataclasses import is_dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from huggingface_hub import HfApi, get_token, run_job
|
||||
|
||||
from .annotate import build_pod_setup
|
||||
from .dataset import ensure_dataset_available
|
||||
from .hf import _pod_forwarded_args, follow_job, resolve_job_tags
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.annotations.camera_curation.config import CameraCurationConfig
|
||||
|
||||
# Same rationale as the annotate submitter: --root is host-local, --repo_id is
|
||||
# re-emitted, config files can't be read on the pod, and --job could smuggle a
|
||||
# remote target back onto the pod.
|
||||
_SUBMITTER_OWNED_ARGS = ("--root", "--repo_id", "--config_path", "--job")
|
||||
|
||||
|
||||
def _local_config_file_args(cfg: CameraCurationConfig) -> list[str]:
|
||||
return ["--config_path", *(f"--{name}" for name in vars(cfg) if is_dataclass(getattr(cfg, name)))]
|
||||
|
||||
|
||||
def build_pod_command(repo_id: str, lerobot_ref: str, argv: list[str]) -> list[str]:
|
||||
"""``bash -c`` command the pod runs: setup prelude, then curate-cameras."""
|
||||
forwarded = _pod_forwarded_args(argv, drop_names=_SUBMITTER_OWNED_ARGS, drop_prefixes=("--job.",))
|
||||
curate = shlex.join(
|
||||
["lerobot-curate-cameras", f"--repo_id={repo_id}", *forwarded, "--job.target=local"]
|
||||
)
|
||||
return ["bash", "-c", f"{build_pod_setup(lerobot_ref)} && {curate}"]
|
||||
|
||||
|
||||
def submit_curate_to_hf(cfg: CameraCurationConfig) -> None:
|
||||
"""Submit a camera-curation run to HF Jobs and tail its logs."""
|
||||
token = get_token()
|
||||
if not token:
|
||||
raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.")
|
||||
|
||||
if cfg.repo_id is None:
|
||||
raise ValueError(
|
||||
"Remote curation requires --repo_id: the pod downloads the dataset from the Hub, "
|
||||
"and --root only names a directory on this machine."
|
||||
)
|
||||
|
||||
argv = sys.argv[1:]
|
||||
passed = {tok.split("=", 1)[0] for tok in argv}
|
||||
used_config_files = sorted(passed.intersection(_local_config_file_args(cfg)))
|
||||
if used_config_files:
|
||||
raise ValueError(
|
||||
f"{', '.join(used_config_files)} cannot be used with a remote --job.target: the pod "
|
||||
"cannot read config files from this machine. Pass the settings as CLI flags instead."
|
||||
)
|
||||
|
||||
if cfg.mode == "report":
|
||||
print(
|
||||
"WARNING: --mode=report writes its result into the pod's local copy, which is discarded "
|
||||
"when the job ends. Use --mode=rename to commit the result to the Hub."
|
||||
)
|
||||
|
||||
api = HfApi(token=token)
|
||||
tags = resolve_job_tags(cfg.job.tags)
|
||||
ensure_dataset_available(cfg.repo_id, api=api, tags=tags)
|
||||
|
||||
command = build_pod_command(cfg.repo_id, cfg.job.lerobot_ref, argv)
|
||||
|
||||
print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...")
|
||||
job_info = run_job(
|
||||
image=cfg.job.image,
|
||||
command=command,
|
||||
flavor=cfg.job.target,
|
||||
secrets={"HF_TOKEN": token},
|
||||
timeout=cfg.job.timeout,
|
||||
labels=dict.fromkeys(tags, "true"),
|
||||
)
|
||||
job_id = job_info.id
|
||||
job_url = getattr(job_info, "url", None)
|
||||
print(f"Job submitted: {job_id}")
|
||||
if job_url:
|
||||
print(f" Job page: {job_url}")
|
||||
print(f" Dataset repo: https://huggingface.co/datasets/{cfg.repo_id}")
|
||||
print(f" Monitor: hf jobs logs {job_id}")
|
||||
print(f" Cancel: hf jobs cancel {job_id}")
|
||||
|
||||
if not follow_job(job_id, detach=cfg.job.detach):
|
||||
return
|
||||
|
||||
print("\nCuration complete.")
|
||||
@@ -0,0 +1,244 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""``lerobot-curate-cameras`` — VLM camera-view curation for a LeRobot dataset.
|
||||
|
||||
Downloads only the first episode, asks a VLM to (1) flag blurry/unusable views
|
||||
and (2) label each view (``top``/``wrist``/``front``/…), then either records the
|
||||
result in ``meta/`` (``--mode=report``) or renames the camera keys to
|
||||
``observation.images.<label>`` (``--mode=rename``). For video datasets the
|
||||
rename is a download-free, server-side Hub commit.
|
||||
|
||||
Examples:
|
||||
|
||||
# Cheap, mutation-free triage (writes meta/camera_curation.json):
|
||||
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=report
|
||||
|
||||
# Apply the labels by renaming camera keys on a new branch (video datasets):
|
||||
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=rename --branch=curated
|
||||
|
||||
# Run the VLM decision on a GPU via HF Jobs:
|
||||
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=rename --job.target=h200
|
||||
"""
|
||||
|
||||
import logging
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from lerobot.annotations.camera_curation import curator
|
||||
from lerobot.annotations.camera_curation.config import CameraCurationConfig
|
||||
from lerobot.annotations.steerable_pipeline.frames import make_frame_provider
|
||||
from lerobot.annotations.steerable_pipeline.reader import iter_episodes
|
||||
from lerobot.annotations.steerable_pipeline.vlm_client import make_vlm_client
|
||||
from lerobot.configs import parser
|
||||
from lerobot.utils.constants import HF_LEROBOT_HOME
|
||||
from lerobot.utils.import_utils import _datasets_available, require_package
|
||||
|
||||
if TYPE_CHECKING or _datasets_available:
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _resolve_root(cfg: CameraCurationConfig) -> Path:
|
||||
"""Concrete, writable root for the dataset (never the symlinked snapshot cache)."""
|
||||
if cfg.root is not None:
|
||||
return Path(cfg.root)
|
||||
if cfg.repo_id is not None:
|
||||
return HF_LEROBOT_HOME / cfg.repo_id
|
||||
raise ValueError("Either --repo_id or --root must be provided.")
|
||||
|
||||
|
||||
def _uniform_indices(n: int, k: int) -> list[int]:
|
||||
if n <= 0 or k <= 0:
|
||||
return []
|
||||
if k >= n:
|
||||
return list(range(n))
|
||||
step = (n - 1) / (k - 1) if k > 1 else 0.0
|
||||
return sorted({round(i * step) for i in range(k)})
|
||||
|
||||
|
||||
def _to_uint8_frame(frame: Any) -> Any:
|
||||
"""Scale a float [0,1] image tensor to uint8; pass uint8/PIL through."""
|
||||
import torch
|
||||
|
||||
if isinstance(frame, torch.Tensor) and torch.is_floating_point(frame):
|
||||
return (frame.clamp(0, 1) * 255).to(torch.uint8)
|
||||
return frame
|
||||
|
||||
|
||||
def _sample_frames(dataset: "LeRobotDataset", cfg: CameraCurationConfig) -> dict[str, list[Any]]:
|
||||
"""Sample ``n_frames`` from the inspected episode for each (non-depth) camera.
|
||||
|
||||
Video cameras go through the annotation frame provider (uint8 frames); image
|
||||
cameras are read straight from the dataset rows and scaled to uint8.
|
||||
"""
|
||||
meta = dataset.meta
|
||||
depth_keys = set(meta.depth_keys)
|
||||
video_keys = set(meta.video_keys)
|
||||
image_keys = set(meta.image_keys)
|
||||
cameras = [k for k in meta.camera_keys if k not in depth_keys]
|
||||
|
||||
frames: dict[str, list[Any]] = {k: [] for k in cameras}
|
||||
|
||||
video_cameras = [k for k in cameras if k in video_keys]
|
||||
if video_cameras:
|
||||
provider = make_frame_provider(dataset.root, video_backend=cfg.video_backend)
|
||||
records = list(iter_episodes(dataset.root, only_episodes=(cfg.episode_index,)))
|
||||
record = records[0] if records else None
|
||||
if record is not None:
|
||||
for key in video_cameras:
|
||||
frames[key] = provider.video_for_episode(record, cfg.n_frames, camera_key=key)
|
||||
|
||||
image_cameras = [k for k in cameras if k in image_keys]
|
||||
if image_cameras:
|
||||
n = len(dataset)
|
||||
for i in _uniform_indices(n, cfg.n_frames):
|
||||
item = dataset[i]
|
||||
for key in image_cameras:
|
||||
if key in item:
|
||||
frames[key].append(_to_uint8_frame(item[key]))
|
||||
|
||||
return frames
|
||||
|
||||
|
||||
def _apply_rename(
|
||||
root: Path,
|
||||
dataset: "LeRobotDataset",
|
||||
cfg: CameraCurationConfig,
|
||||
mapping: dict[str, str],
|
||||
verdicts: list["curator.CameraVerdict"],
|
||||
) -> None:
|
||||
"""Apply the computed ``{old: new}`` camera-key mapping.
|
||||
|
||||
Video datasets on the Hub → download-free server-side rename commit.
|
||||
Otherwise (image dataset, local-only, or a swap/cycle) → local
|
||||
``rename_features`` over a full copy of the dataset.
|
||||
"""
|
||||
video_keys = set(dataset.meta.video_keys)
|
||||
all_video = set(mapping) <= video_keys
|
||||
has_swap = bool(set(mapping.values()) & set(mapping))
|
||||
|
||||
if cfg.repo_id is not None and all_video and not has_swap:
|
||||
if cfg.drop_unusable:
|
||||
logger.warning(
|
||||
"--drop_unusable is only applied via the local rename path; the Hub rename keeps "
|
||||
"flagged views (they are still recorded in meta/). Re-run with a local --root to drop."
|
||||
)
|
||||
# Edit a throwaway copy of meta/ so the local cached copy stays pristine
|
||||
# and only the intended files land in the commit.
|
||||
work = Path(tempfile.mkdtemp(prefix="lerobot_curate_"))
|
||||
try:
|
||||
shutil.copytree(root / "meta", work / "meta")
|
||||
commit = curator.rename_camera_keys_on_hub(
|
||||
cfg.repo_id,
|
||||
mapping,
|
||||
work,
|
||||
branch=cfg.branch,
|
||||
commit_message=cfg.push_commit_message,
|
||||
)
|
||||
oid = getattr(commit, "oid", None)
|
||||
ref = cfg.branch or "main"
|
||||
logger.info("Hub rename committed to %s@%s (%s)", cfg.repo_id, ref, oid)
|
||||
finally:
|
||||
shutil.rmtree(work, ignore_errors=True)
|
||||
return
|
||||
|
||||
# Local path: needs the full dataset, so re-load without the episode filter.
|
||||
require_package("datasets", "dataset")
|
||||
from lerobot.datasets import LeRobotDataset as _LeRobotDataset
|
||||
from lerobot.datasets import remove_feature, rename_features
|
||||
|
||||
logger.info("Local rename path (image/local/swap): loading the full dataset from %s", root)
|
||||
full = _LeRobotDataset(cfg.repo_id or "local", root=root)
|
||||
renamed = rename_features(
|
||||
full, mapping, output_dir=root, repo_id=full.repo_id, on_collision=cfg.on_collision
|
||||
)
|
||||
|
||||
if cfg.drop_unusable:
|
||||
unusable_new_keys = [
|
||||
mapping[v.camera_key] for v in verdicts if not v.usable and v.camera_key in mapping
|
||||
]
|
||||
if unusable_new_keys:
|
||||
logger.info("Dropping unusable views: %s", unusable_new_keys)
|
||||
renamed = remove_feature(renamed, unusable_new_keys, output_dir=root, repo_id=renamed.repo_id)
|
||||
|
||||
if cfg.push_to_hub:
|
||||
logger.info("Pushing renamed dataset to %s", renamed.repo_id)
|
||||
renamed.push_to_hub()
|
||||
|
||||
|
||||
@parser.wrap()
|
||||
def curate_cameras(cfg: CameraCurationConfig) -> None:
|
||||
"""Run the camera-view curation pipeline over a dataset's first episode."""
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
|
||||
if cfg.mode not in ("report", "rename"):
|
||||
raise ValueError(f"--mode must be 'report' or 'rename', got {cfg.mode!r}")
|
||||
|
||||
if cfg.job.is_remote:
|
||||
from lerobot.jobs.curate import submit_curate_to_hf
|
||||
|
||||
return submit_curate_to_hf(cfg)
|
||||
|
||||
require_package("datasets", "dataset")
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
root = _resolve_root(cfg)
|
||||
logger.info("curate-cameras: repo_id=%s root=%s mode=%s", cfg.repo_id, root, cfg.mode)
|
||||
|
||||
# Only episode ``cfg.episode_index`` is fetched (a cheap partial download).
|
||||
dataset = LeRobotDataset(
|
||||
cfg.repo_id or "local",
|
||||
root=root,
|
||||
episodes=[cfg.episode_index],
|
||||
download_videos=True,
|
||||
)
|
||||
|
||||
frames = _sample_frames(dataset, cfg)
|
||||
n_with_frames = sum(1 for v in frames.values() if v)
|
||||
logger.info("curate-cameras: %d camera(s), %d with sampled frames", len(frames), n_with_frames)
|
||||
|
||||
vlm = make_vlm_client(cfg.vlm)
|
||||
verdicts = curator.curate_cameras(frames, cfg, vlm)
|
||||
for v in verdicts:
|
||||
logger.info(
|
||||
" %s -> label=%s usable=%s%s",
|
||||
v.camera_key,
|
||||
v.view_label,
|
||||
v.usable,
|
||||
"" if v.usable else f" (blur_reason={v.blur_reason!r})",
|
||||
)
|
||||
|
||||
mapping = curator.build_name_mapping(verdicts, dataset.meta.features, cfg)
|
||||
report_path = curator.write_report(dataset.root, verdicts, mapping, cfg)
|
||||
logger.info("curate-cameras: report written to %s", report_path)
|
||||
logger.info("curate-cameras: proposed rename mapping: %s", mapping or "(none)")
|
||||
|
||||
if cfg.mode == "rename":
|
||||
if not mapping:
|
||||
logger.info("curate-cameras: nothing to rename (no confident labels differ from current keys)")
|
||||
return
|
||||
_apply_rename(dataset.root, dataset, cfg, mapping, verdicts)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
curate_cameras()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -108,6 +108,12 @@ Remove camera feature:
|
||||
--operation.type remove_feature \
|
||||
--operation.feature_names "['observation.image']"
|
||||
|
||||
Rename features/camera keys (no pixel data is re-encoded):
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht \
|
||||
--operation.type rename_features \
|
||||
--operation.name_mapping '{"observation.images.cam_0": "observation.images.left_wrist"}'
|
||||
|
||||
Modify tasks - set a single task for all episodes (WARNING: modifies in-place):
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht \
|
||||
@@ -257,6 +263,7 @@ from lerobot.datasets import (
|
||||
recompute_stats,
|
||||
reencode_dataset,
|
||||
remove_feature,
|
||||
rename_features,
|
||||
split_dataset,
|
||||
)
|
||||
from lerobot.utils.constants import HF_LEROBOT_HOME
|
||||
@@ -298,6 +305,16 @@ class RemoveFeatureConfig(OperationConfig):
|
||||
feature_names: list[str] | None = None
|
||||
|
||||
|
||||
@OperationConfig.register_subclass("rename_features")
|
||||
@dataclass
|
||||
class RenameFeaturesConfig(OperationConfig):
|
||||
# Mapping of {old_feature_key: new_feature_key}, e.g.
|
||||
# {"observation.images.cam_0": "observation.images.left_wrist"}.
|
||||
name_mapping: dict[str, str] | None = None
|
||||
# "error" raises on colliding targets; "suffix" disambiguates (top -> top_2).
|
||||
on_collision: str = "error"
|
||||
|
||||
|
||||
@OperationConfig.register_subclass("modify_tasks")
|
||||
@dataclass
|
||||
class ModifyTasksConfig(OperationConfig):
|
||||
@@ -545,6 +562,42 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
|
||||
LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
|
||||
|
||||
|
||||
def handle_rename_features(cfg: EditDatasetConfig) -> None:
|
||||
if not isinstance(cfg.operation, RenameFeaturesConfig):
|
||||
raise ValueError("Operation config must be RenameFeaturesConfig")
|
||||
|
||||
if not cfg.operation.name_mapping:
|
||||
raise ValueError("name_mapping must be specified for rename_features operation")
|
||||
|
||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||
output_repo_id, output_dir = get_output_path(
|
||||
cfg.repo_id,
|
||||
new_repo_id=cfg.new_repo_id,
|
||||
root=cfg.root,
|
||||
new_root=cfg.new_root,
|
||||
)
|
||||
|
||||
# In case of in-place modification, make the dataset point to the backup directory
|
||||
if output_dir == dataset.root:
|
||||
dataset.root = dataset.root.with_name(dataset.root.name + "_old")
|
||||
|
||||
logging.info(f"Renaming features {cfg.operation.name_mapping} in {cfg.repo_id}")
|
||||
new_dataset = rename_features(
|
||||
dataset,
|
||||
name_mapping=cfg.operation.name_mapping,
|
||||
output_dir=output_dir,
|
||||
repo_id=output_repo_id,
|
||||
on_collision=cfg.operation.on_collision,
|
||||
)
|
||||
|
||||
logging.info(f"Dataset saved to {output_dir}")
|
||||
logging.info(f"Features: {list(new_dataset.meta.features.keys())}")
|
||||
|
||||
if cfg.push_to_hub:
|
||||
logging.info(f"Pushing to hub as {output_repo_id}")
|
||||
LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
|
||||
|
||||
|
||||
def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
|
||||
if not isinstance(cfg.operation, ModifyTasksConfig):
|
||||
raise ValueError("Operation config must be ModifyTasksConfig")
|
||||
@@ -830,6 +883,8 @@ def edit_dataset(cfg: EditDatasetConfig) -> None:
|
||||
handle_merge(cfg)
|
||||
elif operation_type == "remove_feature":
|
||||
handle_remove_feature(cfg)
|
||||
elif operation_type == "rename_features":
|
||||
handle_rename_features(cfg)
|
||||
elif operation_type == "modify_tasks":
|
||||
handle_modify_tasks(cfg)
|
||||
elif operation_type == "convert_image_to_video":
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for the camera-view curation pipeline (stubbed VLM, mocked Hub)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import PIL.Image
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
|
||||
|
||||
import pandas as pd # noqa: E402
|
||||
|
||||
from lerobot.annotations.camera_curation.config import CameraCurationConfig # noqa: E402
|
||||
from lerobot.annotations.camera_curation.curator import ( # noqa: E402
|
||||
CameraVerdict,
|
||||
build_name_mapping,
|
||||
curate_cameras,
|
||||
is_valid_view_label,
|
||||
rename_camera_keys_on_hub,
|
||||
write_report,
|
||||
)
|
||||
from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient # noqa: E402
|
||||
from lerobot.datasets.io_utils import load_info, write_info # noqa: E402
|
||||
from lerobot.datasets.utils import DatasetInfo # noqa: E402
|
||||
from lerobot.utils.io_utils import load_json, write_json # noqa: E402
|
||||
|
||||
VOCAB = ("top", "wrist", "front", "bottom", "left", "right")
|
||||
|
||||
|
||||
def _queued_vlm(responses: list) -> StubVlmClient:
|
||||
"""Stub VLM that returns queued responses in batch order."""
|
||||
state = {"i": 0}
|
||||
|
||||
def responder(_messages):
|
||||
r = responses[state["i"]]
|
||||
state["i"] += 1
|
||||
return r
|
||||
|
||||
return StubVlmClient(responder=responder)
|
||||
|
||||
|
||||
def _tiny_image() -> PIL.Image.Image:
|
||||
return PIL.Image.new("RGB", (16, 12))
|
||||
|
||||
|
||||
def _make_min_meta(root: Path, camera_key: str, dtype: str = "video") -> None:
|
||||
"""Write a minimal ``meta/`` tree with one camera + one action feature."""
|
||||
(root / "meta" / "episodes" / "chunk-000").mkdir(parents=True, exist_ok=True)
|
||||
features = {
|
||||
camera_key: {
|
||||
"dtype": dtype,
|
||||
"shape": (64, 96, 3),
|
||||
"names": ["height", "width", "channels"],
|
||||
"info": {"video.fps": 10.0} if dtype == "video" else None,
|
||||
},
|
||||
"action": {"dtype": "float32", "shape": (2,), "names": None},
|
||||
}
|
||||
write_info(DatasetInfo(codebase_version="v3.0", fps=10, features=features), root)
|
||||
write_json({camera_key: {"mean": [0.0]}, "action": {"mean": [0.0]}}, root / "meta" / "stats.json")
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
"episode_index": [0],
|
||||
f"videos/{camera_key}/from_timestamp": [0.0],
|
||||
f"videos/{camera_key}/to_timestamp": [1.0],
|
||||
f"videos/{camera_key}/chunk_index": [0],
|
||||
f"videos/{camera_key}/file_index": [0],
|
||||
f"stats/{camera_key}/mean": [[0.0]],
|
||||
}
|
||||
)
|
||||
df.to_parquet(root / "meta" / "episodes" / "chunk-000" / "file-000.parquet")
|
||||
|
||||
|
||||
# ------------------------------ pure logic ------------------------------
|
||||
|
||||
|
||||
def test_is_valid_view_label():
|
||||
assert is_valid_view_label("top", VOCAB, allow_combos=True)
|
||||
assert is_valid_view_label("left_wrist", VOCAB, allow_combos=True)
|
||||
assert not is_valid_view_label("left_wrist", VOCAB, allow_combos=False)
|
||||
assert not is_valid_view_label("banana", VOCAB, allow_combos=True)
|
||||
assert not is_valid_view_label("left_left", VOCAB, allow_combos=True) # duplicate token
|
||||
assert not is_valid_view_label("top_wrist_front", VOCAB, allow_combos=True) # 3 tokens
|
||||
assert not is_valid_view_label("", VOCAB, allow_combos=True)
|
||||
|
||||
|
||||
def test_curate_cameras_parses_and_validates(tmp_path):
|
||||
cfg = CameraCurationConfig(view_vocabulary=VOCAB)
|
||||
frames = {
|
||||
"observation.images.a": [_tiny_image()],
|
||||
"observation.images.b": [_tiny_image()],
|
||||
"observation.images.c": [], # no frames -> reported, not sent to the VLM
|
||||
}
|
||||
vlm = _queued_vlm(
|
||||
[
|
||||
{"usable": True, "blur_reason": None, "view_label": "Left Wrist", "confidence": 0.9},
|
||||
{"usable": False, "blur_reason": "out of focus", "view_label": "banana", "confidence": 0.2},
|
||||
]
|
||||
)
|
||||
verdicts = {v.camera_key: v for v in curate_cameras(frames, cfg, vlm)}
|
||||
|
||||
assert verdicts["observation.images.a"].view_label == "left_wrist" # normalized
|
||||
assert verdicts["observation.images.a"].usable is True
|
||||
assert verdicts["observation.images.b"].usable is False
|
||||
assert verdicts["observation.images.b"].blur_reason == "out of focus"
|
||||
assert verdicts["observation.images.b"].view_label is None # invalid label dropped
|
||||
assert verdicts["observation.images.c"].view_label is None # no frames
|
||||
|
||||
|
||||
def test_build_name_mapping_and_collision():
|
||||
cfg = CameraCurationConfig(view_vocabulary=VOCAB)
|
||||
existing = {"observation.images.cam_0": {}, "observation.images.cam_1": {}, "observation.images.top": {}}
|
||||
verdicts = [
|
||||
CameraVerdict("observation.images.cam_0", usable=True, view_label="left_wrist"),
|
||||
CameraVerdict("observation.images.cam_1", usable=True, view_label="front"),
|
||||
# already canonical -> skipped by build_name_mapping
|
||||
CameraVerdict("observation.images.top", usable=True, view_label="top"),
|
||||
]
|
||||
mapping = build_name_mapping(verdicts, existing, cfg)
|
||||
assert mapping == {
|
||||
"observation.images.cam_0": "observation.images.left_wrist",
|
||||
"observation.images.cam_1": "observation.images.front",
|
||||
}
|
||||
# proposed_new_key stamped back onto the verdicts
|
||||
assert verdicts[0].proposed_new_key == "observation.images.left_wrist"
|
||||
|
||||
# two cameras wanting the same label collide under the default policy
|
||||
clash = [
|
||||
CameraVerdict("observation.images.cam_0", usable=True, view_label="top"),
|
||||
CameraVerdict("observation.images.cam_1", usable=True, view_label="top"),
|
||||
]
|
||||
with pytest.raises(ValueError, match="collision"):
|
||||
build_name_mapping(clash, {"observation.images.cam_0": {}, "observation.images.cam_1": {}}, cfg)
|
||||
|
||||
|
||||
def test_write_report(tmp_path):
|
||||
_make_min_meta(tmp_path, "observation.images.cam_0", dtype="video")
|
||||
cfg = CameraCurationConfig(repo_id="user/ds", view_vocabulary=VOCAB)
|
||||
verdicts = [
|
||||
CameraVerdict("observation.images.cam_0", usable=True, view_label="left_wrist", confidence=0.9)
|
||||
]
|
||||
mapping = {"observation.images.cam_0": "observation.images.left_wrist"}
|
||||
|
||||
report_path = write_report(tmp_path, verdicts, mapping, cfg)
|
||||
|
||||
report = load_json(report_path)
|
||||
cam = report["cameras"]["observation.images.cam_0"]
|
||||
assert cam["view_label"] == "left_wrist"
|
||||
assert cam["proposed_new_key"] == "observation.images.left_wrist"
|
||||
# verdict stamped into info.json so it travels with the dataset
|
||||
info = load_info(tmp_path)
|
||||
assert info.features["observation.images.cam_0"]["info"]["curation"]["view_label"] == "left_wrist"
|
||||
|
||||
|
||||
# ------------------------- lightweight Hub rename -------------------------
|
||||
|
||||
|
||||
def test_rename_camera_keys_on_hub_builds_ops(tmp_path):
|
||||
from huggingface_hub import CommitOperationAdd, CommitOperationCopy, CommitOperationDelete
|
||||
|
||||
camera_key = "observation.images.cam_0"
|
||||
new_key = "observation.images.left_wrist"
|
||||
_make_min_meta(tmp_path, camera_key, dtype="video")
|
||||
old_mp4 = f"videos/{camera_key}/chunk-000/file-000.mp4"
|
||||
|
||||
fake_api = MagicMock()
|
||||
fake_api.list_repo_files.return_value = [old_mp4, "meta/info.json", "data/chunk-000/file-000.parquet"]
|
||||
fake_api.create_commit.return_value = MagicMock(oid="deadbeef")
|
||||
|
||||
with patch("huggingface_hub.HfApi", return_value=fake_api):
|
||||
rename_camera_keys_on_hub("user/ds", {camera_key: new_key}, tmp_path, branch="curated")
|
||||
|
||||
kwargs = fake_api.create_commit.call_args.kwargs
|
||||
ops = kwargs["operations"]
|
||||
copies = [o for o in ops if isinstance(o, CommitOperationCopy)]
|
||||
deletes = [o for o in ops if isinstance(o, CommitOperationDelete)]
|
||||
adds = [o for o in ops if isinstance(o, CommitOperationAdd)]
|
||||
|
||||
new_mp4 = f"videos/{new_key}/chunk-000/file-000.mp4"
|
||||
assert any(o.src_path_in_repo == old_mp4 and o.path_in_repo == new_mp4 for o in copies)
|
||||
assert any(o.path_in_repo == old_mp4 for o in deletes)
|
||||
assert any(o.path_in_repo == "meta/info.json" for o in adds)
|
||||
assert kwargs["revision"] == "curated"
|
||||
|
||||
# meta on disk was actually remapped
|
||||
info = load_info(tmp_path)
|
||||
assert new_key in info.features and camera_key not in info.features
|
||||
|
||||
|
||||
def test_rename_camera_keys_on_hub_rejects_image_keys(tmp_path):
|
||||
camera_key = "observation.images.cam_0"
|
||||
_make_min_meta(tmp_path, camera_key, dtype="image")
|
||||
with patch("huggingface_hub.HfApi", return_value=MagicMock()):
|
||||
with pytest.raises(NotImplementedError, match="image data"):
|
||||
rename_camera_keys_on_hub("user/ds", {camera_key: "observation.images.top"}, tmp_path)
|
||||
|
||||
|
||||
def test_rename_camera_keys_on_hub_rejects_swaps(tmp_path):
|
||||
_make_min_meta(tmp_path, "observation.images.a", dtype="video")
|
||||
with patch("huggingface_hub.HfApi", return_value=MagicMock()):
|
||||
with pytest.raises(NotImplementedError, match="swap"):
|
||||
rename_camera_keys_on_hub(
|
||||
"user/ds",
|
||||
{
|
||||
"observation.images.a": "observation.images.b",
|
||||
"observation.images.b": "observation.images.a",
|
||||
},
|
||||
tmp_path,
|
||||
)
|
||||
@@ -23,6 +23,7 @@ import torch
|
||||
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
import pandas as pd # noqa: E402
|
||||
|
||||
from lerobot.configs import DepthEncoderConfig, RGBEncoderConfig
|
||||
from lerobot.datasets.dataset_tools import (
|
||||
@@ -34,9 +35,11 @@ from lerobot.datasets.dataset_tools import (
|
||||
modify_tasks,
|
||||
reencode_dataset,
|
||||
remove_feature,
|
||||
rename_features,
|
||||
split_dataset,
|
||||
)
|
||||
from lerobot.datasets.io_utils import load_info
|
||||
from lerobot.datasets.dataset_tools import _resolve_rename_collisions
|
||||
from lerobot.datasets.io_utils import load_info, load_stats
|
||||
from tests.datasets.test_video_encoding import require_h264, require_hevc, require_libsvtav1
|
||||
from tests.fixtures.constants import DUMMY_DEPTH_FEATURES, DUMMY_DEPTH_KEY
|
||||
from tests.fixtures.dataset_factories import add_frames
|
||||
@@ -1492,3 +1495,162 @@ def test_reencode_dataset_multi_key_multiprocessing(
|
||||
for vk in dataset.meta.video_keys:
|
||||
persisted_encoder = RGBEncoderConfig.from_video_info(persisted_info.features[vk].get("info", {}))
|
||||
assert persisted_encoder == target_cfg
|
||||
|
||||
|
||||
# ----------------------------- rename_features -----------------------------
|
||||
|
||||
|
||||
def _mock_hub(tmp_path):
|
||||
"""Context managers that stop dataset reload from hitting the Hub."""
|
||||
return (
|
||||
patch("lerobot.datasets.dataset_metadata.get_safe_version", return_value="v3.0"),
|
||||
patch(
|
||||
"lerobot.datasets.dataset_metadata.snapshot_download",
|
||||
side_effect=lambda repo_id, **kwargs: str(kwargs.get("local_dir", tmp_path)),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def two_camera_image_dataset(tmp_path, empty_lerobot_dataset_factory):
|
||||
"""An image dataset with two camera views (for collision tests)."""
|
||||
features = {
|
||||
"action": {"dtype": "float32", "shape": (6,), "names": None},
|
||||
"observation.images.cam_0": {"dtype": "image", "shape": (32, 32, 3), "names": None},
|
||||
"observation.images.cam_1": {"dtype": "image", "shape": (32, 32, 3), "names": None},
|
||||
}
|
||||
dataset = empty_lerobot_dataset_factory(root=tmp_path / "two_cam", features=features)
|
||||
for _ in range(2):
|
||||
for _ in range(4):
|
||||
dataset.add_frame(
|
||||
{
|
||||
"action": np.random.randn(6).astype(np.float32),
|
||||
"observation.images.cam_0": np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8),
|
||||
"observation.images.cam_1": np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8),
|
||||
"task": "t",
|
||||
}
|
||||
)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
return dataset
|
||||
|
||||
|
||||
def test_resolve_rename_collisions_error_and_suffix():
|
||||
features = {"a": {}, "b": {}, "c": {}}
|
||||
# many-to-one
|
||||
with pytest.raises(ValueError, match="same target"):
|
||||
_resolve_rename_collisions({"a": "top", "b": "top"}, features, "error")
|
||||
# target collides with an untouched key
|
||||
with pytest.raises(ValueError, match="existing feature"):
|
||||
_resolve_rename_collisions({"a": "c"}, features, "error")
|
||||
# suffix disambiguates deterministically
|
||||
resolved = _resolve_rename_collisions({"a": "top", "b": "top"}, features, "suffix")
|
||||
assert set(resolved.values()) == {"top", "top_2"}
|
||||
assert resolved["a"] == "top" # sorted-source order keeps the first
|
||||
|
||||
|
||||
def test_rename_image_feature(sample_dataset, tmp_path):
|
||||
old, new = "observation.images.top", "observation.images.wrist"
|
||||
m1, m2 = _mock_hub(tmp_path)
|
||||
with m1, m2:
|
||||
renamed = rename_features(sample_dataset, {old: new}, output_dir=tmp_path / "renamed")
|
||||
|
||||
assert new in renamed.meta.features
|
||||
assert old not in renamed.meta.features
|
||||
assert renamed.meta.features[new]["dtype"] == "image"
|
||||
# the frame still decodes under the new key
|
||||
item = renamed[0]
|
||||
assert new in item and old not in item
|
||||
# stats moved to the new key
|
||||
stats = load_stats(renamed.root)
|
||||
assert new in stats and old not in stats
|
||||
|
||||
|
||||
@require_h264
|
||||
def test_rename_video_feature_no_reencode(tmp_path, empty_lerobot_dataset_factory, features_factory):
|
||||
features = features_factory(use_videos=True) # observation.images.{laptop,phone}
|
||||
dataset = empty_lerobot_dataset_factory(root=tmp_path / "vid", features=features, use_videos=True)
|
||||
add_frames(dataset, num_frames=4)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
old, new = "laptop", "observation.images.top" # features_factory uses bare camera keys
|
||||
old_bytes = (dataset.root / dataset.meta.get_video_file_path(0, old)).read_bytes()
|
||||
|
||||
m1, m2 = _mock_hub(tmp_path)
|
||||
with m1, m2:
|
||||
renamed = rename_features(dataset, {old: new}, output_dir=tmp_path / "renamed")
|
||||
|
||||
assert new in renamed.meta.features and old not in renamed.meta.features
|
||||
new_mp4 = renamed.root / renamed.meta.get_video_file_path(0, new)
|
||||
assert new_mp4.exists()
|
||||
# a rename must not re-encode: the mp4 is byte-identical.
|
||||
assert new_mp4.read_bytes() == old_bytes
|
||||
# episodes metadata columns were remapped.
|
||||
ep_parquet = next((renamed.root / "meta" / "episodes").glob("*/*.parquet"))
|
||||
cols = pd.read_parquet(ep_parquet).columns
|
||||
assert f"videos/{new}/from_timestamp" in cols
|
||||
assert f"videos/{old}/from_timestamp" not in cols
|
||||
# the video still decodes under the new key.
|
||||
assert new in renamed[0]
|
||||
|
||||
|
||||
def test_rename_collision_raises(two_camera_image_dataset, tmp_path):
|
||||
m1, m2 = _mock_hub(tmp_path)
|
||||
with m1, m2, pytest.raises(ValueError, match="collision"):
|
||||
rename_features(
|
||||
two_camera_image_dataset,
|
||||
{
|
||||
"observation.images.cam_0": "observation.images.top",
|
||||
"observation.images.cam_1": "observation.images.top",
|
||||
},
|
||||
output_dir=tmp_path / "out",
|
||||
)
|
||||
|
||||
|
||||
def test_rename_collision_suffix(two_camera_image_dataset, tmp_path):
|
||||
m1, m2 = _mock_hub(tmp_path)
|
||||
with m1, m2:
|
||||
renamed = rename_features(
|
||||
two_camera_image_dataset,
|
||||
{
|
||||
"observation.images.cam_0": "observation.images.top",
|
||||
"observation.images.cam_1": "observation.images.top",
|
||||
},
|
||||
output_dir=tmp_path / "out",
|
||||
on_collision="suffix",
|
||||
)
|
||||
keys = set(renamed.meta.features)
|
||||
assert {"observation.images.top", "observation.images.top_2"} <= keys
|
||||
|
||||
|
||||
def test_rename_identity_only_raises(sample_dataset, tmp_path):
|
||||
with pytest.raises(ValueError, match="identity"):
|
||||
rename_features(
|
||||
sample_dataset,
|
||||
{"observation.images.top": "observation.images.top"},
|
||||
output_dir=tmp_path / "out",
|
||||
)
|
||||
|
||||
|
||||
def test_rename_missing_key_raises(sample_dataset, tmp_path):
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
rename_features(
|
||||
sample_dataset,
|
||||
{"observation.images.nope": "observation.images.top"},
|
||||
output_dir=tmp_path / "out",
|
||||
)
|
||||
|
||||
|
||||
def test_rename_required_feature_raises(sample_dataset, tmp_path):
|
||||
with pytest.raises(ValueError, match="required"):
|
||||
rename_features(sample_dataset, {"timestamp": "t2"}, output_dir=tmp_path / "out")
|
||||
|
||||
|
||||
def test_rename_slash_in_target_raises(sample_dataset, tmp_path):
|
||||
with pytest.raises(ValueError, match="'/'"):
|
||||
rename_features(
|
||||
sample_dataset,
|
||||
{"observation.images.top": "observation/images/top"},
|
||||
output_dir=tmp_path / "out",
|
||||
)
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# The script imports ``lerobot.datasets`` (via the annotation frame provider),
|
||||
# which only ships under the ``dataset`` extra.
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
from lerobot.scripts.lerobot_curate_cameras import _to_uint8_frame, _uniform_indices # noqa: E402
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"n,k,expected",
|
||||
[
|
||||
(0, 4, []),
|
||||
(5, 0, []),
|
||||
(3, 5, [0, 1, 2]), # k >= n -> all frames
|
||||
(1, 4, [0]),
|
||||
(10, 1, [0]),
|
||||
(10, 4, [0, 3, 6, 9]), # evenly spaced, endpoints included
|
||||
],
|
||||
)
|
||||
def test_uniform_indices(n, k, expected):
|
||||
assert _uniform_indices(n, k) == expected
|
||||
|
||||
|
||||
def test_to_uint8_frame_scales_floats():
|
||||
frame = torch.ones(3, 4, 4, dtype=torch.float32) # [0,1] float
|
||||
out = _to_uint8_frame(frame)
|
||||
assert out.dtype == torch.uint8
|
||||
assert int(out.max()) == 255
|
||||
|
||||
|
||||
def test_to_uint8_frame_passthrough_uint8():
|
||||
frame = torch.zeros(3, 4, 4, dtype=torch.uint8)
|
||||
out = _to_uint8_frame(frame)
|
||||
assert out is frame # uint8 passes through untouched
|
||||
Reference in New Issue
Block a user