#!/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, )