feat(dataset filtering): adding support for VLM based dataset filtering following lerobot annotation pipeline style

tests/annotations/test_camera_curation.py

@@ -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,
+            )

tests/datasets/test_dataset_tools.py

@@ -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",
+        )

tests/scripts/test_lerobot_curate_cameras.py

@@ -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