Spaces:
Running
Running
File size: 7,415 Bytes
d646dd7 af470e1 d646dd7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | """Filesystem artifact discovery shared by review APIs.
The active application reads only the canonical project-level
``data/meta/videos`` layout. Old sidecar directories are intentionally not
resolved by the runtime.
"""
from __future__ import annotations
import re
from pathlib import Path
_SOURCE_SUFFIXES = ("_depth_rgb", "_depth", "_rgb")
def _episode_index_for_id(session_dir: Path, episode_id: str) -> int | None:
try:
from app.project_dataset import episode_row
row = episode_row(Path(session_dir), str(episode_id))
return int(row.get("episode_index")) if row else None
except (TypeError, ValueError, OSError):
return None
def _safe(value: str) -> str:
return str(value or "").replace("/", "_").replace("\\", "_").strip()
def _source_variants(value: str) -> set[str]:
raw = _safe(value).lower()
if not raw:
return set()
variants = {raw, re.sub(r"[^a-z0-9_.-]+", "_", raw)}
changed = True
while changed:
changed = False
for item in tuple(variants):
for suffix in _SOURCE_SUFFIXES:
if item.endswith(suffix) and len(item) > len(suffix):
shortened = item[: -len(suffix)].rstrip("_")
if shortened not in variants:
variants.add(shortened)
changed = True
# Collectors have used both ``rgb`` and ``head_rgb`` for the
# same head-facing colour stream. Keep this alias limited to
# RGB keys; depth keys must never become interchangeable with an
# RGB stream merely because they share a device prefix.
if item == "head_rgb":
variants.add("rgb")
elif item == "rgb":
variants.add("head_rgb")
elif item.endswith("_head_rgb"):
variants.add(item[:-len("_head_rgb")] + "_rgb")
elif item.endswith("_rgb") and not item.endswith("_depth_rgb"):
variants.add(item[:-len("_rgb")] + "_head_rgb")
return variants
def source_matches(value: str | None, camera: str) -> bool:
"""Whether a manifest/file source refers to *camera*.
``D405`` and ``D405_depth_rgb`` intentionally match: the former is the
keypoint artifact stem while the latter is the uploaded video source key.
"""
return bool(_source_variants(str(value or "")) & _source_variants(camera))
def _unique(paths: list[Path]) -> list[Path]:
seen: set[str] = set()
result: list[Path] = []
for path in paths:
key = str(path)
if key not in seen:
seen.add(key)
result.append(path)
return result
def find_hand_keypoint_candidates(session_dir: Path, camera: str,
episode_id: str | None = None) -> list[Path]:
"""Find MediaPipe 2D parquet files for a camera."""
result: list[Path] = []
# Current runs merge the detector columns into the canonical episode
# parquet. Returning that file keeps the review API independent of a
# per-run output directory.
canonical = _canonical_episode_data(session_dir, episode_id, "2d")
if canonical is not None:
result.append(canonical)
return _unique(result)
def find_hand3d_candidates(session_dir: Path, camera: str,
episode_id: str | None = None) -> list[Path]:
"""Find 3D artifacts whose manifest/source maps to a camera."""
result: list[Path] = []
canonical = _canonical_episode_data(session_dir, episode_id, "3d")
if canonical is not None:
result.append(canonical)
return _unique(result)
def _canonical_episode_data(session_dir: Path, episode_id: str | None,
dimension: str) -> Path | None:
"""Return the canonical episode parquet when it contains merged results."""
if not episode_id:
return None
try:
from app.project_dataset import episode_row, episode_chunk_for_index, is_project_dataset
root = Path(session_dir)
if not is_project_dataset(root):
return None
row = episode_row(root, str(episode_id))
if row is None:
return None
index = int(row.get("episode_index", 0))
path = (root / "data" / f"chunk-{episode_chunk_for_index(index):03d}"
/ f"episode_{index:06d}.parquet")
if not path.is_file():
return None
try:
import pyarrow.parquet as pq
# Use logical top-level names; physical schema names flatten nested
# landmark lists to ``element`` and make valid 3D artifacts look
# empty.
columns = {str(name) for name in
pq.ParquetFile(path).schema_arrow.names}
except Exception:
import pandas as pd
columns = {str(name) for name in
pd.read_parquet(path, columns=[], engine="pyarrow").columns}
if dimension == "3d":
marker = any(
"landmarks_3d" in name or "world_position" in name
or "hand_3d" in name
or name.endswith("hand_left_3d")
or name.endswith("hand_right_3d")
for name in columns)
else:
marker = any("keypoints" in name or "2d_present" in name
or "hand_0_present" in name for name in columns)
return path if marker else None
except Exception:
return None
def find_depth_video(session_dir: Path, source: str,
episode_id: str | None = None) -> Path | None:
# LeRobot v2.1 is authoritative. Active project datasets must resolve
# depth only from their canonical metric video; stale PNG/raw-depth
# sidecars are deliberately not allowed to become a preview source.
try:
from app.lerobot_v21 import is_depth_source, iter_video_streams
episode_index = _episode_index_for_id(Path(session_dir), str(episode_id))
for source_key, path in iter_video_streams(Path(session_dir) / "videos"):
if (is_depth_source(source_key)
and source_matches(source_key, source)
and (episode_index is None
or f"episode_{episode_index:06d}" in path.stem)):
return path
except Exception:
pass
return None
def has_depth_source(session_dir: Path, episode_id: str | None = None) -> bool:
"""Whether an episode has a usable canonical depth video.
PNG sequences can still be found by migration utilities, but they do not
qualify as an active project's depth input or world-coordinate preview.
"""
try:
from app.lerobot_v21 import is_depth_source, iter_video_streams
from app.project_dataset import episode_row
wanted_index = None
if episode_id and (Path(session_dir) / "meta" / "episodes").is_dir():
row = episode_row(Path(session_dir), str(episode_id))
wanted_index = int(row["episode_index"]) if row else None
for source, path in iter_video_streams(Path(session_dir) / "videos"):
if not is_depth_source(source):
continue
if wanted_index is None or path.stem == f"episode_{wanted_index:06d}":
return True
except (TypeError, ValueError, OSError):
pass
return False
|