Spaces:
Running
Running
| """Clean up active LeRobot v2.1 media without touching recoverable history. | |
| The first D435 collectors used several names for the same two physical | |
| streams. This utility normalizes them to ``D435_rgb`` and ``D435_depth`` in | |
| active datasets. It also removes files advertised as pure depth when ffprobe | |
| shows that they are colour/video pixels rather than metric | |
| ``gray12le``/``gray16le`` depth. Every moved video and every rewritten | |
| metadata file is preserved in ``data/.backups/session-migrations``. | |
| Default mode is a read-only report. Use ``--apply`` only after reviewing it. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import re | |
| import shutil | |
| import subprocess | |
| import sys | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import Any | |
| APP_ROOT = Path(__file__).resolve().parents[1] | |
| if str(APP_ROOT) not in sys.path: | |
| sys.path.insert(0, str(APP_ROOT)) | |
| from app.lerobot_v21 import ( # noqa: E402 | |
| VIDEO_SUFFIXES, | |
| canonical_source_key, | |
| is_depth_source, | |
| iter_video_streams, | |
| ) | |
| from app.project_dataset import ( # noqa: E402 | |
| episode_files, | |
| episode_row, | |
| project_episode_rows, | |
| write_project_episode_index, | |
| ) | |
| EPISODE_RE = re.compile(r"episode_(\d+)$") | |
| SOURCE_RENAMES = { | |
| "D435_depth_rgb": "D435_rgb", | |
| "D435_head_rgb": "D435_rgb", | |
| "D435_head_depth": "D435_depth", | |
| } | |
| DEPTH_PIX_FMTS = {"gray12le", "gray16le"} | |
| DEVICE_NAME_ALIASES = { | |
| "d435_depth": "D435", | |
| "d435_head": "D435", | |
| } | |
| def _probe(path: Path) -> dict[str, str]: | |
| try: | |
| result = subprocess.run( | |
| ["ffprobe", "-v", "error", "-select_streams", "v:0", | |
| "-show_entries", "stream=codec_name,pix_fmt,width,height", | |
| "-of", "json", str(path)], | |
| capture_output=True, text=True, timeout=30, check=False, | |
| ) | |
| streams = json.loads(result.stdout or "{}").get("streams") or [] | |
| return {str(k): str(v or "") for k, v in (streams[0] if streams else {}).items()} | |
| except (OSError, ValueError, TypeError, subprocess.TimeoutExpired): | |
| return {} | |
| def _copy_or_link(source: Path, target: Path) -> None: | |
| target.parent.mkdir(parents=True, exist_ok=True) | |
| if target.exists(): | |
| return | |
| try: | |
| os.link(source, target) | |
| except OSError: | |
| shutil.copy2(source, target) | |
| def _backup_file(source: Path, project: Path, backup: Path) -> None: | |
| _copy_or_link(source, backup / project.name / source.relative_to(project)) | |
| def _backup_video(source: Path, project: Path, backup: Path) -> None: | |
| _copy_or_link(source, backup / project.name / source.relative_to(project)) | |
| def _metadata_files(project: Path) -> list[Path]: | |
| meta = project / "meta" | |
| if not meta.is_dir(): | |
| return [] | |
| return sorted(path for path in meta.rglob("*") | |
| if path.is_file() and path.suffix.lower() in {".json", ".jsonl"}) | |
| def _merge_values(current: Any, incoming: Any) -> Any: | |
| """Merge colliding source-key metadata while keeping canonical values.""" | |
| if isinstance(current, dict) and isinstance(incoming, dict): | |
| merged = dict(current) | |
| for key, value in incoming.items(): | |
| if key not in merged: | |
| merged[key] = value | |
| else: | |
| merged[key] = _merge_values(merged[key], value) | |
| return merged | |
| if isinstance(current, list) and isinstance(incoming, list): | |
| result = list(current) | |
| for value in incoming: | |
| if value not in result: | |
| result.append(value) | |
| return result | |
| return current | |
| def _rename_source_value(value: Any, old: str, new: str) -> Any: | |
| old_feature = f"observation.images.{old}" | |
| new_feature = f"observation.images.{new}" | |
| if isinstance(value, str): | |
| return value.replace(old_feature, new_feature).replace(old, new) | |
| if isinstance(value, list): | |
| return [_rename_source_value(item, old, new) for item in value] | |
| if not isinstance(value, dict): | |
| return value | |
| result: dict[str, Any] = {} | |
| for key, child in value.items(): | |
| renamed_key = key.replace(old_feature, new_feature).replace(old, new) | |
| renamed_child = _rename_source_value(child, old, new) | |
| if renamed_key in result: | |
| result[renamed_key] = _merge_values(result[renamed_key], renamed_child) | |
| else: | |
| result[renamed_key] = renamed_child | |
| return result | |
| def _rewrite_source_names(project: Path, old: str, new: str, | |
| backup: Path | None) -> list[str]: | |
| changed: list[str] = [] | |
| for path in _metadata_files(project): | |
| try: | |
| text = path.read_text(encoding="utf-8") | |
| except (OSError, UnicodeError): | |
| continue | |
| updated = text | |
| if path.suffix.lower() == ".json": | |
| try: | |
| value = json.loads(text) | |
| updated = json.dumps( | |
| _rename_source_value(value, old, new), | |
| ensure_ascii=False, indent=2, | |
| ) | |
| except (ValueError, TypeError): | |
| updated = text.replace( | |
| f"observation.images.{old}", | |
| f"observation.images.{new}", | |
| ).replace(old, new) | |
| elif path.suffix.lower() == ".jsonl": | |
| lines = [] | |
| for line in text.splitlines(): | |
| if not line.strip(): | |
| continue | |
| try: | |
| lines.append(json.dumps( | |
| _rename_source_value(json.loads(line), old, new), | |
| ensure_ascii=False, | |
| )) | |
| except (ValueError, TypeError): | |
| lines.append(line.replace( | |
| f"observation.images.{old}", | |
| f"observation.images.{new}", | |
| ).replace(old, new)) | |
| updated = "\n".join(lines) + ("\n" if lines else "") | |
| else: | |
| updated = text.replace( | |
| f"observation.images.{old}", | |
| f"observation.images.{new}", | |
| ).replace(old, new) | |
| if updated == text: | |
| continue | |
| if backup is not None: | |
| _backup_file(path, project, backup) | |
| path.write_text(updated, encoding="utf-8") | |
| changed.append(str(path.relative_to(project))) | |
| return changed | |
| def _canonical_device_name(value: Any, kind: Any = "", key: Any = "") -> Any: | |
| if not isinstance(value, str): | |
| return value | |
| low = value.strip().lower() | |
| context = f"{kind} {key}".lower() | |
| if low in DEVICE_NAME_ALIASES and "d435" in context: | |
| return DEVICE_NAME_ALIASES[low] | |
| return value | |
| def _normalise_devices(items: Any) -> Any: | |
| if not isinstance(items, list): | |
| return items | |
| merged: list[Any] = [] | |
| by_key: dict[str, int] = {} | |
| for raw in items: | |
| if not isinstance(raw, dict): | |
| merged.append(raw) | |
| continue | |
| item = _normalise_device_metadata(raw) | |
| key = str(item.get("key") or "").strip() | |
| if not key or key not in by_key: | |
| if key: | |
| by_key[key] = len(merged) | |
| merged.append(item) | |
| continue | |
| index = by_key[key] | |
| merged[index] = _merge_values(merged[index], item) | |
| return merged | |
| def _normalise_device_metadata(value: Any, field: str = "", | |
| kind: Any = "", key: Any = "") -> Any: | |
| """Normalize physical D435 names without changing stream keys.""" | |
| if isinstance(value, list): | |
| return [_normalise_device_metadata(item, field, kind, key) | |
| for item in value] | |
| if not isinstance(value, dict): | |
| if field in {"name", "device", "device_name"}: | |
| return _canonical_device_name(value, kind, key) | |
| return value | |
| local_kind = value.get("kind") or kind | |
| local_key = value.get("key") or key | |
| result: dict[str, Any] = {} | |
| for name, child in value.items(): | |
| if name == "devices": | |
| child = _normalise_devices(child) | |
| elif name == "device_names" and isinstance(child, dict): | |
| child = { | |
| stream: _canonical_device_name(display, "d435", stream) | |
| for stream, display in child.items() | |
| } | |
| elif name in {"name", "device", "device_name"}: | |
| child = _canonical_device_name(child, local_kind, local_key) | |
| else: | |
| child = _normalise_device_metadata(child, name, local_kind, local_key) | |
| result[name] = child | |
| return result | |
| def _rewrite_device_metadata(project: Path, backup: Path | None) -> list[str]: | |
| changed: list[str] = [] | |
| for path in _metadata_files(project): | |
| if path.suffix.lower() != ".json": | |
| continue | |
| try: | |
| before = path.read_text(encoding="utf-8") | |
| value = json.loads(before) | |
| except (OSError, UnicodeError, ValueError, TypeError): | |
| continue | |
| updated_value = _normalise_device_metadata(value) | |
| updated = json.dumps(updated_value, ensure_ascii=False, indent=2) | |
| if updated == before: | |
| continue | |
| if backup is not None: | |
| _backup_file(path, project, backup) | |
| path.write_text(updated, encoding="utf-8") | |
| changed.append(str(path.relative_to(project))) | |
| return changed | |
| def _rewrite_video_extensions(project: Path, backup: Path | None) -> list[str]: | |
| """Make the v2.1 extension index agree with active video files.""" | |
| info_path = project / "meta" / "info.json" | |
| if not info_path.is_file(): | |
| return [] | |
| actual = { | |
| source: path.suffix.lower().lstrip(".") | |
| for source, path in iter_video_streams(project / "videos") | |
| } | |
| if not actual: | |
| return [] | |
| try: | |
| before = info_path.read_text(encoding="utf-8") | |
| info = json.loads(before) | |
| except (OSError, UnicodeError, ValueError, TypeError): | |
| return [] | |
| if not isinstance(info, dict): | |
| return [] | |
| current = info.get("video_extensions") | |
| if current == actual: | |
| return [] | |
| info["video_extensions"] = dict(sorted(actual.items())) | |
| if backup is not None: | |
| _backup_file(info_path, project, backup) | |
| info_path.write_text(json.dumps(info, ensure_ascii=False, indent=2), | |
| encoding="utf-8") | |
| return [str(info_path.relative_to(project))] | |
| def _rewrite_json(path: Path, transform) -> bool: | |
| try: | |
| value = json.loads(path.read_text(encoding="utf-8")) | |
| except (OSError, UnicodeError, ValueError, TypeError): | |
| return False | |
| updated = transform(value) | |
| if updated == value: | |
| return False | |
| temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") | |
| temporary.write_text(json.dumps(updated, ensure_ascii=False, indent=2), | |
| encoding="utf-8") | |
| os.replace(temporary, path) | |
| return True | |
| def _remove_source_from_value(value: Any, source: str) -> Any: | |
| """Remove a source only from source-bearing metadata fields.""" | |
| source_fields = {"cameras", "device_names"} | |
| list_fields = {"slots", "depth_sources", "metric_depth_sources", "camera_sources"} | |
| if isinstance(value, list): | |
| return [_remove_source_from_value(item, source) for item in value | |
| if item != source] | |
| if not isinstance(value, dict): | |
| return value | |
| result: dict[str, Any] = {} | |
| for key, item in value.items(): | |
| if key in source_fields and isinstance(item, dict): | |
| item = {name: child for name, child in item.items() if name != source} | |
| elif key in list_fields and isinstance(item, list): | |
| item = [child for child in item if child != source] | |
| result[key] = _remove_source_from_value(item, source) | |
| return result | |
| def _remove_episode_source_metadata(project: Path, episode_id: str, | |
| source: str, backup: Path | None) -> list[str]: | |
| changed: list[str] = [] | |
| collector = project / "meta" / "collector" / episode_id | |
| if not collector.is_dir(): | |
| return changed | |
| for path in sorted(collector.glob("*.json")): | |
| before = path.read_text(encoding="utf-8", errors="ignore") | |
| if source not in before: | |
| continue | |
| if backup is not None: | |
| _backup_file(path, project, backup) | |
| if _rewrite_json(path, lambda value: _remove_source_from_value(value, source)): | |
| changed.append(str(path.relative_to(project))) | |
| return changed | |
| def _remove_episode_row_source(project: Path, episode_index: int, source: str, | |
| backup: Path | None) -> bool: | |
| rows = project_episode_rows(project) | |
| if not rows: | |
| return False | |
| changed = False | |
| prefix = f"videos/observation.images.{source}/" | |
| updated: list[dict[str, Any]] = [] | |
| for row in rows: | |
| if int(row.get("episode_index", -1)) == episode_index: | |
| filtered = {key: value for key, value in row.items() | |
| if not str(key).startswith(prefix)} | |
| changed |= filtered != row | |
| row = filtered | |
| updated.append(row) | |
| if not changed: | |
| return False | |
| if backup is not None: | |
| for path in sorted((project / "meta" / "episodes").rglob("*.parquet")): | |
| _backup_file(path, project, backup) | |
| write_project_episode_index(project, updated) | |
| return True | |
| def _episode_id(project: Path, index: int) -> str: | |
| for row in project_episode_rows(project): | |
| if int(row.get("episode_index", -1)) == index: | |
| return str(row.get("episode_id") or row.get("source_batch") or "") | |
| return "" | |
| def _active_video_paths(project: Path) -> list[Path]: | |
| videos = project / "videos" | |
| if not videos.is_dir(): | |
| return [] | |
| return sorted(path for path in videos.rglob("*.mp4") | |
| if path.is_file() and any(part.startswith("chunk-") for part in path.parts)) | |
| def clean_project(project: Path, backup: Path | None, apply: bool) -> dict[str, Any]: | |
| result: dict[str, Any] = {"project": project.name, "renamed": [], | |
| "invalid_depth": [], "metadata": []} | |
| if not project.is_dir(): | |
| return result | |
| # Normalize historical D435 source names in every chunk. If the target | |
| # directory already exists, merge only non-colliding episode files into | |
| # it; the caller has already verified that episode numbers do not overlap. | |
| for old_source, new_source in SOURCE_RENAMES.items(): | |
| old_dirs = sorted((project / "videos").rglob( | |
| f"observation.images.{old_source}")) | |
| for old_dir in old_dirs: | |
| if not old_dir.is_dir(): | |
| continue | |
| new_dir = old_dir.with_name(f"observation.images.{new_source}") | |
| result["renamed"].append({ | |
| "from": str(old_dir.relative_to(project)), | |
| "to": str(new_dir.relative_to(project)), | |
| }) | |
| if not apply: | |
| continue | |
| if new_dir.exists(): | |
| conflicts = sorted( | |
| item.name for item in old_dir.iterdir() | |
| if (new_dir / item.name).exists() | |
| ) | |
| if conflicts: | |
| raise RuntimeError( | |
| f"Rename collision: {new_dir} ({', '.join(conflicts)})" | |
| ) | |
| if backup is not None: | |
| for path in old_dir.rglob("*"): | |
| if path.is_file(): | |
| _backup_video(path, project, backup) | |
| new_dir.mkdir(parents=True, exist_ok=True) | |
| for item in sorted(old_dir.iterdir()): | |
| target = new_dir / item.name | |
| if target.exists(): | |
| raise RuntimeError(f"Rename collision: {target}") | |
| item.rename(target) | |
| old_dir.rmdir() | |
| if apply and old_dirs: | |
| result["metadata"].extend(_rewrite_source_names( | |
| project, old_source, new_source, backup)) | |
| # Merge duplicate physical-device records left by historical uploads. | |
| # Stream keys remain untouched: D435_rgb and D435_depth are data sources, | |
| # while D435 is the single hardware name shown by the UI. | |
| if apply: | |
| result["metadata"].extend(_rewrite_device_metadata(project, backup)) | |
| result["metadata"].extend(_rewrite_video_extensions(project, backup)) | |
| # A metric depth stream must be lossless gray12/gray16. Keep the valid | |
| # files and move only the colour-coded/ordinary-video members out of the | |
| # active dataset. This handles mixed historical batches in one source. | |
| invalid: list[tuple[Path, str, int]] = [] | |
| for path in _active_video_paths(project): | |
| relative_parts = path.relative_to(project).parts | |
| source_parts = [part for part in relative_parts | |
| if part.startswith("observation.images.")] | |
| if not source_parts: | |
| continue | |
| source = canonical_source_key(source_parts[-1]) | |
| if not is_depth_source(source): | |
| continue | |
| probe = _probe(path) | |
| if probe.get("pix_fmt") in DEPTH_PIX_FMTS: | |
| continue | |
| match = EPISODE_RE.search(path.stem) | |
| if not match: | |
| continue | |
| invalid.append((path, source, int(match.group(1)))) | |
| for path, source, episode_index in invalid: | |
| result["invalid_depth"].append({ | |
| "path": str(path.relative_to(project)), | |
| "source": source, | |
| "episode_index": episode_index, | |
| "probe": _probe(path), | |
| }) | |
| if not apply: | |
| continue | |
| if backup is not None: | |
| _backup_video(path, project, backup) | |
| path.unlink() | |
| episode_id = _episode_id(project, episode_index) | |
| if episode_id: | |
| result["metadata"].extend(_remove_episode_source_metadata( | |
| project, episode_id, source, backup)) | |
| if _remove_episode_row_source(project, episode_index, source, backup): | |
| result["metadata"].append("meta/episodes/") | |
| return result | |
| def clean_duplicate_processed_depth(root: Path, backup: Path | None, | |
| apply: bool) -> list[dict[str, Any]]: | |
| """Move raw PNG sidecars when the same episode has canonical depth video.""" | |
| processed = root / "processed" | |
| sessions = root / "sessions" | |
| result: list[dict[str, Any]] = [] | |
| if not processed.is_dir(): | |
| return result | |
| for raw_root in sorted(processed.rglob("raw_depth")): | |
| if not raw_root.is_dir(): | |
| continue | |
| try: | |
| relative = raw_root.relative_to(processed) | |
| project_name, episode_name = relative.parts[0], relative.parts[1] | |
| except (ValueError, IndexError): | |
| continue | |
| project = sessions / project_name | |
| if not project.is_dir(): | |
| continue | |
| row = episode_row(project, episode_name) | |
| if not row: | |
| continue | |
| try: | |
| episode_index = int(row["episode_index"]) | |
| except (KeyError, TypeError, ValueError): | |
| continue | |
| canonical = { | |
| source: path for source, path in episode_files(project, episode_index)["videos"] | |
| } | |
| for source_dir in sorted(path for path in raw_root.rglob("*") if path.is_dir()): | |
| pngs = [path for path in source_dir.glob("*.png") | |
| if path.is_file() and path.stem.isdigit()] | |
| if not pngs: | |
| continue | |
| source = canonical_source_key(source_dir.name) | |
| video = canonical.get(source) | |
| if video is None or not is_depth_source(source): | |
| continue | |
| probe = _probe(video) | |
| if probe.get("pix_fmt") not in DEPTH_PIX_FMTS: | |
| continue | |
| item = { | |
| "path": str(source_dir.relative_to(root)), | |
| "project": project_name, | |
| "episode_id": episode_name, | |
| "episode_index": episode_index, | |
| "source": source, | |
| "png_count": len(pngs), | |
| "canonical_video": str(video.relative_to(root)), | |
| "canonical_probe": probe, | |
| } | |
| result.append(item) | |
| if apply: | |
| if backup is None: | |
| raise RuntimeError("backup path is required for cleanup") | |
| target = backup / "processed-legacy" / source_dir.relative_to(root) | |
| target.parent.mkdir(parents=True, exist_ok=True) | |
| original_parent = source_dir.parent | |
| shutil.move(str(source_dir), str(target)) | |
| # Remove only the empty legacy raw_depth shell left behind in | |
| # the active processed run; never remove a non-empty run or | |
| # any sibling annotation output. | |
| parent = original_parent | |
| while parent != processed and parent != root: | |
| try: | |
| parent.rmdir() | |
| except OSError: | |
| break | |
| parent = parent.parent | |
| return result | |
| def main() -> int: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--root", type=Path, required=True, | |
| help="active data root containing sessions/") | |
| parser.add_argument("--project", action="append", | |
| help="limit to one or more project directory names") | |
| parser.add_argument("--apply", action="store_true", | |
| help="perform moves and metadata updates; default is dry-run") | |
| args = parser.parse_args() | |
| root = args.root.expanduser().resolve() | |
| sessions = root / "sessions" | |
| if not sessions.is_dir(): | |
| raise SystemExit(f"sessions directory not found: {sessions}") | |
| wanted = set(args.project or []) | |
| projects = [path for path in sorted(sessions.iterdir()) | |
| if path.is_dir() and (not wanted or path.name in wanted)] | |
| stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") | |
| backup = root / ".backups" / "session-migrations" / f"media-cleanup-{stamp}" \ | |
| if args.apply else None | |
| if backup is not None: | |
| backup.mkdir(parents=True, exist_ok=False) | |
| report = {"root": str(root), "apply": bool(args.apply), | |
| "backup": str(backup) if backup else None, "projects": [], | |
| "legacy_processed_depth": []} | |
| for project in projects: | |
| report["projects"].append(clean_project(project, backup, args.apply)) | |
| report["legacy_processed_depth"] = clean_duplicate_processed_depth( | |
| root, backup, args.apply) | |
| if args.apply: | |
| for item in report["projects"]: | |
| project = sessions / item["project"] | |
| if backup is None: | |
| continue | |
| # Keep maintenance reports outside the canonical LeRobot meta | |
| # namespace. ``meta`` should contain dataset metadata only. | |
| manifest = backup / "reports" / f"{item['project']}.json" | |
| manifest.parent.mkdir(parents=True, exist_ok=True) | |
| manifest.write_text(json.dumps(item, ensure_ascii=False, indent=2), | |
| encoding="utf-8") | |
| print(json.dumps(report, ensure_ascii=False, indent=2)) | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |