Blame for processor/app/lerobot_v21.py

"""LeRobot v2.1 storage helpers used by the upload pipeline.

The collector currently sends a small v3-style archive for one episode.  The
server normalizes each upload into the project-level ``sessions/<project>``
dataset and appends it to the v2.1 layout:

    data/chunk-000/episode_XXXXXX.parquet
    meta/{info.json,stats.json,tasks.json,episodes/chunk-*/episode_*.parquet,...}
    videos/observation.images.<source>/chunk-000/episode_XXXXXX.mp4

Workflow results are merged into the corresponding episode parquet files
under the canonical ``data/`` and ``meta/`` directories.  Preview videos are
rendered in the browser and are never stored in the dataset.
"""

from __future__ import annotations

import json
import math
import re
import shutil
import subprocess
from uuid import uuid4
from pathlib import Path
from typing import Any, Iterable


VIDEO_SUFFIXES = {".mp4", ".mkv", ".mov", ".avi", ".m4v", ".webm"}
DEPTH_PNG_SUFFIX = ".png"

# Canonical depth-video contract.  The stored gray12le samples are logarithmic
# depth codes, never JET/RGB pixels and never direct millimetres.
DEPTH_MIN_MM = 100.0
DEPTH_MAX_MM = 5000.0
DEPTH_QMAX = 4095
DEPTH_QP = 6
DEPTH_VIDEO_ENCODING = "depth_mm_log_to_gray12le"
# Kept as aliases for callers that imported the old names.
DEPTH_VIDEO_MAX_MM = DEPTH_QMAX
DEPTH_PREVIEW_MIN_MM = DEPTH_MIN_MM
DEPTH_PREVIEW_MAX_MM = DEPTH_MAX_MM
EPISODES_METADATA_CHUNK_SIZE = 1000

_LOG_LO = math.log(DEPTH_MIN_MM)
_LOG_SPAN = math.log(DEPTH_MAX_MM) - _LOG_LO
_LOG_STEP = DEPTH_QMAX / _LOG_SPAN

# Historical collectors used ``D435_depth_rgb`` for the colour stream of a
# D435.  It is not a depth stream (``is_depth_source`` already excludes it),
# but keeping that name in the canonical dataset makes the RGB/depth split
# ambiguous.  New uploads and legacy-path discovery use the clear key below.
SOURCE_ALIASES = {
    "d435_depth_rgb": "D435_rgb",
    "d435_head_rgb": "D435_rgb",
    "d435_head_depth": "D435_depth",
}
DEVICE_NAME_ALIASES = {
    "d435_depth": "D435",
    "d435_head": "D435",
}


def depth_video_encoder_args() -> list[str]:
    """Return the exact FFmpeg arguments for canonical depth video output.

    keyint=2:官方 pyav 解码按"最近关键帧向后 seek"取帧,大 GOP(默认
    250)时 seek 会越过目标帧落在下一个关键帧 → FrameTimestampError
    (实测 870 帧视频仅 4 个关键帧,viz/训练在关键帧前 1-2 帧崩溃)。
    官方录制管线 GOP=2,这里对齐。
    """
    return [
        "-pix_fmt", "gray12le",
        "-tag:v", "hvc1",
        "-x265-params", f"qp={DEPTH_QP}:range=full:keyint=2:min-keyint=2",
    ]


def quantize_depth(depth_mm):
    """Convert uint16 millimetres to the canonical 12-bit log depth code."""
    import numpy as np

    mm = np.asarray(depth_mm, dtype=np.float64)
    valid = mm > 0
    codes = np.zeros(mm.shape, dtype="<u2")
    if valid.any():
        codes[valid] = np.clip(
            np.rint((np.log(mm[valid]) - _LOG_LO) * _LOG_STEP),
            0, DEPTH_QMAX,
        ).astype("<u2")
    return codes


def dequantize_depth(codes):
    """Convert canonical 12-bit log codes to rounded uint16 millimetres."""
    import numpy as np

    c = np.asarray(codes, dtype=np.float64)
    return np.rint(np.exp(_LOG_LO + c / _LOG_STEP)).astype("<u2")


def codes_to_heatmap_bgr(codes):
    """Convert canonical codes to the exact OpenCV JET BGR display image."""
    import cv2
    import numpy as np

    c8 = ((np.clip(codes, 0, DEPTH_QMAX).astype(np.int32) * 255)
          // DEPTH_QMAX).astype(np.uint8)
    return cv2.applyColorMap(c8, cv2.COLORMAP_JET)


def depth_to_heatmap_bgr(depth_mm):
    """Convert millimetres to the canonical preview-only OpenCV JET image."""
    return codes_to_heatmap_bgr(quantize_depth(depth_mm))


# Turbo colormap (Google) polynomial coefficients — matches Rerun's default
# depth visualization (blue → cyan → green → yellow → red continuous ramp).
_TURBO_POLY = {
    "r": [0.13572138, 4.61539260, -42.66032258, 132.13108234,
          -152.94239396, 59.28637943],
    "g": [0.09140261, 2.19418839, 4.84296658, -14.18503333,
          4.27729857, 2.82956604],
    "b": [0.10667330, 12.64194608, -60.58204836, 110.36276771,
          -89.90310912, 27.34824973],
}
_TURBO_LUT: list[list[int]] | None = None


def turbo_lut(size: int = 256) -> list[list[int]]:
    """N-entry Google Turbo colormap (RGB), generated once per process.

    The frontend uses 256 entries; the preview transcode uses a 4096-entry
    cube so the range mapping keeps 12-bit precision (a 256-entry LUT in the
    8-bit domain caused visible banding on smooth depth gradients).
    """
    global _TURBO_LUT
    if size == 256 and _TURBO_LUT is not None:
        return _TURBO_LUT
    lut = []
    for i in range(size):
        x = i / max(1, size - 1)
        rgb = []
        for channel in ("r", "g", "b"):
            c = _TURBO_POLY[channel]
            v = c[0] + x * (c[1] + x * (c[2] + x * (c[3] + x * (c[4] + x * c[5]))))
            rgb.append(int(max(0, min(255, round(v * 255)))))
        lut.append(rgb)
    if size == 256:
        _TURBO_LUT = lut
    return lut


def turbo_lut_256() -> list[list[int]]:
    """256-entry Google Turbo colormap (RGB), generated once per process."""
    return turbo_lut(256)


def _depth_video_info_from_path(path: Path) -> dict[str, Any]:
    """Read depth metadata beside a canonical video when available."""
    path = Path(path)
    source = next(
        (part[len("observation.images."):]
         for part in path.parts if part.startswith("observation.images.")),
        path.parent.name,
    )
    for parent in (path.parent, *path.parents):
        info_path = parent / "meta" / "info.json"
        if not info_path.is_file():
            continue
        info = _read_json(info_path)
        feature = (info.get("features") or {}).get(
            f"observation.images.{source}", {})
        video_info = dict(feature.get("video_info")) if isinstance(feature, dict) and isinstance(feature.get("video_info"), dict) else {}
        # 逐批次覆盖:episodes 索引行可声明 video.depth_encoding,
        # 用于"同一会话混合新旧格式"的历史数据(如 9/4 对数批次与
        # 8 月 mm/10 批次并存于一个会话)。
        try:
            import pyarrow.parquet as pq
            idx = info_path.parent / "episodes" / "chunk-000" / f"{path.stem}.parquet"
            if idx.is_file():
                row = pq.read_table(idx).to_pandas().iloc[0]
                override = row.get("video.depth_encoding")
                if override and str(override).strip():
                    video_info["video.depth_encoding"] = str(override).strip()
        except Exception:
            pass
        return video_info
    return {}


def _depth_video_mode(video_info: dict[str, Any]) -> str:
    """Return ``log``, ``legacy_log``, ``legacy_mm10`` or ``direct_mm`` for a stored stream."""
    if not isinstance(video_info, dict):
        video_info = {}
    encoding = str(video_info.get("video.depth_encoding") or "").lower()
    # 8 月老采集端深度按 0.1 系数存储(毫米/10),显式声明为 legacy_mm10。
    # 该检查必须在 uint16_mm 检查之前,否则会被当成直接毫米读出 10 倍偏小。
    if "mm10" in encoding or "mm/10" in encoding:
        return "legacy_mm10"
    # Older uploads explicitly identify the gray12le samples as direct
    # millimetres.  This check must happen before the generic gray12le
    # fallback below; otherwise a direct-mm stream is interpreted as a log
    # code stream and values such as 500 mm are decoded to the wrong depth.
    if "uint16_mm" in encoding and "gray12le" in encoding:
        return "direct_mm"
    if encoding == DEPTH_VIDEO_ENCODING or "log_to_gray12le" in encoding:
        return "log"
    # The current collector contract stores the logarithmic 12-bit code
    # directly in gray12le.  Some older ``meta/info.json`` files describe the
    # transport as ``uint16_mm_to_gray12le_clipped_4095mm`` even though the
    # pixels are already codes, not millimetres.  A pure depth gray12le
    # stream is therefore log-coded by default; otherwise values such as
    # code=1511 are incorrectly treated as 1511 mm instead of ~420 mm.
    if (video_info.get("video.is_depth_map")
            and "gray12le" in encoding):
        return "log"
    if video_info.get("video.use_log") and (
            video_info.get("video.depth_min") is not None
            and video_info.get("video.depth_max") is not None):
        return "legacy_log"
    # Existing project videos were written as direct millimetres.  Keeping
    # this fallback prevents an old recording from being silently interpreted
    # as a new logarithmic code stream before it is explicitly migrated.
    return "direct_mm"


def _read_json(path: Path) -> dict[str, Any]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError, TypeError):
        return {}
    return value if isinstance(value, dict) else {}


def _write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2),
                    encoding="utf-8")


def _safe_source_key(value: str) -> str:
    value = str(value or "").strip()
    if value.startswith("observation.images."):
        value = value[len("observation.images."):]
    value = re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("._")
    return value or "camera"


def canonical_source_key(value: str) -> str:
    """Return the stable source key used by the v2.1 dataset.

    This is deliberately a narrow compatibility map.  Generic ``*_rgb`` and
    ``*_depth`` names must remain distinct because they describe different
    physical streams and are used by workflow matching.
    """
    safe = _safe_source_key(value)
    return SOURCE_ALIASES.get(safe.lower(), safe)


def canonical_device_name(value: str, kind: str = "", key: str = "") -> str:
    """Return a stable hardware name without changing stream identifiers."""
    text = str(value or "")
    if (text.strip().lower() in DEVICE_NAME_ALIASES
            and "d435" in f"{kind} {key}".lower()):
        return DEVICE_NAME_ALIASES[text.strip().lower()]
    return text


def normalize_metadata_sources(document: Any) -> dict[str, Any]:
    """Normalize source aliases and duplicate physical-device records.

    This runs on each newly uploaded episode before it is merged into the
    project dataset, so an old collector cannot reintroduce ``D435_head_*``
    keys or a second device entry after a one-time migration.
    """
    if not isinstance(document, dict):
        return document
    result = dict(document)
    for field in ("features", "cameras", "device_names", "video_extensions"):
        value = result.get(field)
        if not isinstance(value, dict):
            continue
        mapped: dict[str, Any] = {}
        for raw_key, child in value.items():
            key = str(raw_key)
            prefix = "observation.images."
            source = key[len(prefix):] if field == "features" and key.startswith(prefix) else key
            canonical = canonical_source_key(source)
            new_key = prefix + canonical if field == "features" and key.startswith(prefix) else canonical
            # A canonical entry is authoritative if both old and new names
            # exist in one metadata document.
            if new_key not in mapped or source == canonical:
                mapped[new_key] = child
        result[field] = mapped

    devices = result.get("devices")
    if isinstance(devices, list):
        merged: list[dict[str, Any]] = []
        positions: dict[str, int] = {}
        for raw in devices:
            if not isinstance(raw, dict):
                continue
            item = dict(raw)
            kind = str(item.get("kind") or "")
            device_key = str(item.get("key") or item.get("id") or "")
            if "name" in item:
                item["name"] = canonical_device_name(item["name"], kind, device_key)
            slots = [canonical_source_key(slot) for slot in item.get("slots") or []
                     if str(slot).strip()]
            item["slots"] = list(dict.fromkeys(slots))
            for field in ("resolution", "fps"):
                mapping = item.get(field)
                if isinstance(mapping, dict):
                    item[field] = {
                        canonical_source_key(k): v for k, v in mapping.items()
                    }
            if device_key and device_key in positions:
                current = merged[positions[device_key]]
                current["slots"] = list(dict.fromkeys(
                    [*(current.get("slots") or []), *(item.get("slots") or [])]
                ))
                for field in ("resolution", "fps"):
                    current.setdefault(field, {})
                    current[field].update(item.get(field) or {})
                for field, value in item.items():
                    if field not in {"slots", "resolution", "fps"} and not current.get(field):
                        current[field] = value
                continue
            if device_key:
                positions[device_key] = len(merged)
            merged.append(item)
        result["devices"] = merged

    device_names = result.get("device_names")
    if isinstance(device_names, dict):
        result["device_names"] = {
            canonical_source_key(stream): canonical_device_name(
                display, "d435", str(stream),
            )
            for stream, display in device_names.items()
        }
    cameras = result.get("cameras")
    if isinstance(cameras, dict):
        normalized_cameras = {}
        for stream, value in cameras.items():
            if isinstance(value, dict):
                value = dict(value)
                for field in ("device", "device_name"):
                    if field in value:
                        value[field] = canonical_device_name(
                            value[field], str(value.get("type") or "d435"), str(stream),
                        )
            normalized_cameras[canonical_source_key(stream)] = value
        result["cameras"] = normalized_cameras
    return result


def source_key_from_video(path: Path, videos_root: Path | None = None) -> str:
    """Extract the source key from legacy, v2.1 and collector-v3 paths."""
    path = Path(path)
    if videos_root is None:
        parts = path.parts
        try:
            index = max(i for i, value in enumerate(parts) if value == "videos")
            videos_root = Path(*parts[:index + 1])
        except ValueError:
            videos_root = path.parent

    try:
        relative = path.relative_to(videos_root)
        parts = relative.parts
    except ValueError:
        parts = path.parts

    # v2.1: videos/observation.images.left/chunk-000/episode_000000.mp4
    # v3.0: videos/chunk-000/observation.images.left/file-000.mp4
    for part in parts[:-1]:
        if part.startswith("observation.images."):
            return canonical_source_key(part)

    # Collector source directories may omit the observation.images prefix.
    # Keep parsing them for normalization, but all newly written output uses
    # one of the two canonical layouts above.
    if len(parts) >= 2:
        candidate = parts[-2]
        if candidate.startswith("chunk-") or candidate.startswith("chunk_"):
            candidate = parts[-3] if len(parts) >= 3 else path.stem
        return canonical_source_key(candidate)
    return canonical_source_key(path.stem)


def is_depth_source(source_key: str | None) -> bool:
    """Return true only for a pure depth stream, not ``*_depth_rgb``."""
    value = str(source_key or "").lower().strip()
    if not value or "depth" not in value:
        return False
    return not any(token in value for token in ("rgb", "color", "video"))


def iter_video_streams(videos_root: Path) -> list[tuple[str, Path]]:
    """Return one stable ``(source_key, path)`` entry per video stream.

    The normalizer writes one file per source for a single uploaded batch.  If
    a malformed archive contains multiple files for one source, the caller
    can detect the duplicate instead of silently mixing episodes.
    """
    if not videos_root.is_dir():
        return []
    result: list[tuple[str, Path]] = []
    for path in sorted(videos_root.rglob("*")):
        if path.is_file() and path.suffix.lower() in VIDEO_SUFFIXES:
            result.append((source_key_from_video(path, videos_root), path))
    return result


def _episode_index_from_name(batch_name: str) -> int | None:
    match = re.search(r"_(\d{6})$", str(batch_name))
    return int(match.group(1)) if match else None


def _data_parquets(root: Path) -> list[Path]:
    data_root = root / "data"
    if not data_root.is_dir():
        return []
    return sorted(
        path for path in data_root.rglob("*.parquet")
        if path.is_file() and "meta" not in path.parts
    )


def _load_frame_table(paths: list[Path]):
    try:
        import pandas as pd
    except ImportError as exc:
        raise RuntimeError("pandas is required for LeRobot v2.1 normalization") from exc

    candidates = []
    for path in paths:
        try:
            table = pd.read_parquet(path)
        except Exception as exc:
            print(f"[LeRobot v2.1] skip unreadable parquet {path.name}: {exc}")
            continue
        if "frame_index" in table.columns:
            candidates.append((len(table), path, table))
    if not candidates:
        raise RuntimeError("No data parquet containing frame_index was found")

    # The main frame table is normally the largest one.  Additional parquet
    # files (for example a glove sensor stream) are merged by frame_index.
    candidates.sort(key=lambda item: (item[0], str(item[1])), reverse=True)
    _, primary_path, frame_table = candidates[0]
    for _, extra_path, extra in reversed(candidates[1:]):
        if extra_path == primary_path:
            continue
        columns = [c for c in extra.columns
                   if c != "frame_index" and c not in frame_table.columns]
        if not columns:
            continue
        frame_table = frame_table.merge(
            extra[["frame_index", *columns]], on="frame_index", how="left",
        )
    return frame_table


def _legacy_frame_table(root: Path, info: dict[str, Any], batch_name: str):
    """Build a truthful frame table for older collector archives.

    Some early RGB-only uploads contain ``timestamps.json`` and episode
    summary parquet files, but no frame-level ``data`` parquet.  We preserve
    the available timeline instead of treating a processed hand-keypoint
    parquet as source data.  Missing action/IMU values are deliberately not
    fabricated.
    """
    import pandas as pd

    timestamp_paths = [root / "timestamps.json",
                       root / "meta" / "collector" / "timestamps.json"]
    timestamp_rows: list[dict[str, Any]] = []
    for path in timestamp_paths:
        if not path.is_file():
            continue
        try:
            payload = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, ValueError, TypeError):
            continue
        rows = payload.get("timestamps") if isinstance(payload, dict) else payload
        if isinstance(rows, list):
            timestamp_rows = [row for row in rows if isinstance(row, dict)]
            if timestamp_rows:
                break

    episode_index = _episode_index_from_name(batch_name)
    length = 0
    for path in sorted((root / "meta").rglob("*.parquet")) if (root / "meta").is_dir() else []:
        try:
            table = pd.read_parquet(path)
        except Exception:
            continue
        if "length" not in table.columns or table.empty:
            continue
        row = table.iloc[0]
        try:
            length = int(row.get("length") or 0)
        except (TypeError, ValueError):
            length = 0
        if episode_index is None:
            try:
                value = int(row.get("episode_index"))
                if value >= 0:
                    episode_index = value
            except (TypeError, ValueError):
                pass
        if length:
            break

    if timestamp_rows:
        frame_table = pd.DataFrame(timestamp_rows)
    else:
        if length <= 0:
            # Last-resort count from the actual source video.  This is still
            # source-derived and does not invent any sensor or pose values.
            length = sum(_probe_video(path, float(info.get("fps") or 30)).get("frames", 0)
                         for _source, path in iter_video_streams(root / "videos"))
        if length <= 0:
            raise RuntimeError("No frame-level data or timestamps found")
        fps = float(info.get("fps") or 30.0)
        frame_table = pd.DataFrame({
            "frame_index": list(range(length)),
            "timestamp": [index / fps for index in range(length)],
        })

    if "frame_index" not in frame_table.columns:
        frame_table.insert(0, "frame_index", range(len(frame_table)))
    if episode_index is None:
        episode_index = 0
    frame_table["episode_index"] = int(episode_index)
    return frame_table


def _infer_episode_index(root: Path, batch_name: str, frame_table) -> int:
    if "episode_index" in frame_table.columns and len(frame_table):
        try:
            value = int(frame_table["episode_index"].iloc[0])
            if value >= 0:
                return value
        except (TypeError, ValueError, IndexError):
            pass
    episode_index = _episode_index_from_name(batch_name)
    return episode_index if episode_index is not None else 0


def _probe_video(path: Path, fallback_fps: float = 30.0) -> dict[str, Any]:
    result: dict[str, Any] = {
        "frames": 0, "fps": float(fallback_fps or 30.0),
        "width": 0, "height": 0, "codec": "", "pix_fmt": "",
    }
    try:
        import cv2
        cap = cv2.VideoCapture(str(path))
        result.update({
            "frames": int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0),
            "fps": float(cap.get(cv2.CAP_PROP_FPS) or fallback_fps or 30.0),
            "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0),
            "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0),
        })
        cap.release()
    except Exception:
        pass
    try:
        completed = 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=15, check=False,
        )
        streams = json.loads(completed.stdout or "{}").get("streams") or []
        if streams:
            stream = streams[0]
            result["codec"] = str(stream.get("codec_name") or "")
            result["pix_fmt"] = str(stream.get("pix_fmt") or "")
            result["width"] = int(stream.get("width") or result["width"] or 0)
            result["height"] = int(stream.get("height") or result["height"] or 0)
    except (OSError, ValueError, TypeError, subprocess.TimeoutExpired):
        pass
    return result


def depth_png_directories(root: Path) -> list[tuple[str, Path, list[Path]]]:
    """Find numeric depth-frame directories in a legacy upload tree."""
    root = Path(root)
    if not root.is_dir():
        return []
    result: list[tuple[str, Path, list[Path]]] = []
    seen: set[str] = set()
    for directory in sorted(path for path in root.rglob("*") if path.is_dir()):
        frames = sorted(
            path for path in directory.iterdir()
            if path.is_file() and path.suffix.lower() == DEPTH_PNG_SUFFIX
            and path.stem.isdigit()
        )
        if not frames:
            continue
        source = _safe_source_key(directory.name)
        if source in seen:
            raise RuntimeError(
                f"Multiple depth frame directories use the same source key: {source}"
            )
        numbers = [int(path.stem) for path in frames]
        expected = list(range(numbers[0], numbers[0] + len(numbers)))
        if numbers != expected:
            raise RuntimeError(
                f"Depth frame sequence has gaps for {source}: "
                f"expected consecutive files, got {frames[0].name}..{frames[-1].name}"
            )
        seen.add(source)
        result.append((source, directory, frames))
    return result


def encode_depth_png_sequence(depth_dir: Path, frames: list[Path], destination: Path,
                              fps: float = 30.0) -> dict[str, Any]:
    """Encode uint16 millimetre PNGs using the native depth-video contract.

    The generated stream contains 12-bit logarithmic depth codes, not a
    colourized preview.  Quantization is explicit so FFmpeg cannot silently
    rescale uint16 PNG values.  CQP qp=6 and full range are part of the
    interchange contract; CRF and lossless mode are intentionally avoided.
    """
    import cv2
    import numpy as np

    if not frames:
        raise RuntimeError(f"No depth frames found in {depth_dir}")
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg is required to convert depth PNGs to video")
    destination = Path(destination)
    destination.parent.mkdir(parents=True, exist_ok=True)
    first_number = int(frames[0].stem)
    sample = cv2.imread(str(frames[0]), cv2.IMREAD_UNCHANGED)
    if sample is None or sample.ndim != 2 or sample.dtype != np.uint16:
        raise RuntimeError(
            f"Depth frame must be a single-channel uint16 PNG: {frames[0]}"
        )
    height, width = sample.shape
    temporary = destination.with_name(
        f".{destination.stem}.tmp-{uuid4().hex}{destination.suffix or '.mp4'}"
    )
    process = None
    try:
        process = subprocess.Popen(
            [ffmpeg, "-y", "-loglevel", "error",
             "-f", "rawvideo", "-pix_fmt", "gray12le",
             "-s", f"{width}x{height}", "-r", str(float(fps or 30.0)),
             "-i", "pipe:0", "-frames:v", str(len(frames)),
             "-an", "-c:v", "libx265",
             *depth_video_encoder_args(), "-movflags", "+faststart",
             str(temporary)],
            stdin=subprocess.PIPE, stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
        )
        assert process.stdin is not None
        for frame_path in frames:
            frame = cv2.imread(str(frame_path), cv2.IMREAD_UNCHANGED)
            if frame is None or frame.ndim != 2 or frame.dtype != np.uint16:
                raise RuntimeError(
                    f"Depth frame must be a single-channel uint16 PNG: {frame_path}"
                )
            if frame.shape != sample.shape:
                raise RuntimeError(
                    f"Depth frame resolution mismatch: {frame_path} "
                    f"has {frame.shape}, expected {sample.shape}"
                )
            # Explicitly provide canonical 12-bit log samples. Passing the
            # PNG pattern to FFmpeg would silently rescale uint16 values.
            codes = quantize_depth(frame)
            process.stdin.write(np.ascontiguousarray(codes).tobytes())
        process.stdin.close()
        stderr = process.stderr.read() if process.stderr is not None else b""
        return_code = process.wait(timeout=1800)
        if return_code != 0 or not temporary.is_file() \
                or temporary.stat().st_size <= 0:
            detail = (stderr.decode(errors="replace")
                      or "ffmpeg returned no output").strip()
            raise RuntimeError(f"Depth video encoding failed: {detail[-1000:]}")
        probe = _probe_video(temporary, float(fps or 30.0))
        if probe.get("codec") not in {"hevc", "h265"} \
                or probe.get("pix_fmt") != "gray12le":
            raise RuntimeError(
                "Depth video verification failed: "
                f"codec={probe.get('codec')!r}, pix_fmt={probe.get('pix_fmt')!r}"
            )
        if int(probe.get("width") or 0) <= 0 or int(probe.get("height") or 0) <= 0:
            raise RuntimeError("Depth video verification failed: invalid dimensions")
        temporary.replace(destination)
        probe["frames"] = int(probe.get("frames") or len(frames))
        probe["fps"] = float(probe.get("fps") or fps or 30.0)
        probe["source_frame_start"] = first_number
        probe["source_frame_count"] = len(frames)
        probe["depth_encoding"] = DEPTH_VIDEO_ENCODING
        probe["depth_min_mm"] = DEPTH_MIN_MM
        probe["depth_max_mm"] = DEPTH_MAX_MM
        probe["depth_qmax"] = DEPTH_QMAX
        probe["depth_qp"] = DEPTH_QP
        return probe
    except (BrokenPipeError, OSError, subprocess.TimeoutExpired) as exc:
        if process is not None and process.poll() is None:
            process.kill()
            process.wait()
        raise RuntimeError(f"Depth video encoding failed: {exc}") from exc
    finally:
        if process is not None and process.poll() is None:
            process.kill()
            process.wait()
        temporary.unlink(missing_ok=True)


class DepthVideoReader:
    """Sequentially decode a depth stream as raw codes or metric millimetres."""

    def __init__(self, path: Path):
        self.path = Path(path)
        ffmpeg = shutil.which("ffmpeg")
        if not ffmpeg:
            raise RuntimeError("ffmpeg is required to read metric depth video")
        probe = _probe_video(self.path)
        pix_fmt = str(probe.get("pix_fmt") or "")
        if pix_fmt not in {"gray12le", "gray16le"}:
            raise RuntimeError(
                f"Metric depth video must use gray12le/gray16le, got {pix_fmt!r}"
            )
        self.width = int(probe.get("width") or 0)
        self.height = int(probe.get("height") or 0)
        if self.width <= 0 or self.height <= 0:
            raise RuntimeError(f"Metric depth video has invalid dimensions: {self.path}")
        self.frame_count = int(probe.get("frames") or 0)
        self.fps = float(probe.get("fps") or 30.0)
        self.pix_fmt = pix_fmt
        self.video_info = _depth_video_info_from_path(self.path)
        self.mode = _depth_video_mode(self.video_info)
        self._frame_bytes = self.width * self.height * 2
        self._process = subprocess.Popen(
            [ffmpeg, "-v", "error", "-i", str(self.path),
             "-f", "rawvideo", "-pix_fmt", pix_fmt, "pipe:1"],
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        )

    def read_codes(self):
        """Read one stored little-endian uint16 frame without colorization."""
        import numpy as np

        if self._process.stdout is None:
            return None
        chunks: list[bytes] = []
        remaining = self._frame_bytes
        while remaining:
            data = self._process.stdout.read(remaining)
            if not data:
                return None
            chunks.append(data)
            remaining -= len(data)
        return np.frombuffer(b"".join(chunks), dtype="<u2").reshape(
            self.height, self.width,
        ).copy()

    def read_mm(self):
        """Read one frame as uint16 millimetres for metric processing."""
        import numpy as np

        codes = self.read_codes()
        if codes is None:
            return None
        return self._codes_to_mm(codes)

    def _codes_to_mm(self, codes):
        import numpy as np

        if self.mode == "log":
            depth = dequantize_depth(codes)
            # The canonical display maps code 0 to JET(0), but metric
            # processing must retain its invalid-pixel semantics.
            depth[codes == 0] = 0
            return depth
        if self.mode == "legacy_mm10":
            # 8 月老采集端把毫米深度按 0.1 系数存储(码值 = mm/10),
            # ×10 还原毫米;0 仍保持无效像素语义。
            depth = codes.astype(np.float64) * 10.0
            depth[codes == 0] = 0.0
            return np.rint(np.clip(depth, 0, 65535)).astype("<u2")
        if self.mode == "legacy_log":
            low = float(self.video_info["video.depth_min"])
            high = float(self.video_info["video.depth_max"])
            shift = float(self.video_info.get("video.shift") or 0.0)
            log_low = math.log(low + shift)
            log_span = math.log(high + shift) - log_low
            depth_m = np.exp(log_low + codes.astype(np.float64) / DEPTH_QMAX * log_span) - shift
            depth_m[codes == 0] = 0.0
            return np.rint(np.clip(depth_m * 1000.0, 0, 65535)).astype("<u2")
        return codes

    def read_canonical_codes(self):
        """Read a frame as canonical codes without creating a color image.

        New streams already contain canonical codes.  Known legacy direct-mm
        and legacy-log streams are converted in memory for compatibility; no
        converted frame is written back to the dataset.
        """
        codes = self.read_codes()
        if codes is None:
            return None
        if self.mode == "log":
            import numpy as np
            return np.clip(codes, 0, DEPTH_QMAX).astype("<u2", copy=False)
        return quantize_depth(self._codes_to_mm(codes))

    def read(self):
        """Backward-compatible metric read used by 3D processing modules."""
        return self.read_mm()

    def close(self) -> None:
        process = self._process
        if process.poll() is None:
            # The decoder usually sits blocked on a full stdout pipe once the
            # caller stops reading.  SIGTERM makes FFmpeg try a graceful
            # flush that never completes in that state (measured 5s timeouts
            # per close), so give it a very short grace window and then kill.
            process.terminate()
            try:
                process.wait(timeout=0.3)
            except subprocess.TimeoutExpired:
                process.kill()
                try:
                    process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    pass
        if process.stdout:
            process.stdout.close()
        if process.stderr:
            process.stderr.close()

    def __enter__(self):
        return self

    def __exit__(self, _exc_type, _exc, _tb):
        self.close()


class DepthCodeSeekReader:
    """Sequential code reader positioned at an exact frame via FFmpeg seek.

    The canonical depth contract encodes keyint=2 (an IDR every two frames),
    so input seeking lands one frame before the target at most.  FFmpeg's
    accurate seek drops the lead-in frames itself (verified: ``-ss`` at an
    exact frame timestamp yields that frame as the first output), which turns
    a scrub jump from an O(N) sequential decode into a constant-time seek.

    The pooled window path replaces its sequential reader with one of these
    after a jump, so forward playback continues from the new position without
    re-seeking per window.
    """

    def __init__(self, path: Path, start_frame: int, fps: float = 30.0,
                 pix_fmt: str = "gray12le"):
        import numpy as np

        self.path = Path(path)
        ffmpeg = shutil.which("ffmpeg")
        if not ffmpeg:
            raise RuntimeError("ffmpeg is required to seek depth video")
        probe = _probe_video(self.path)
        self.width = int(probe.get("width") or 0)
        self.height = int(probe.get("height") or 0)
        if self.width <= 0 or self.height <= 0:
            raise RuntimeError(f"Depth video has invalid dimensions: {self.path}")
        self._frame_bytes = self.width * self.height * 2
        # Keep the stream identity so a later scrub jump on the same pool
        # entry re-derives the seek timestamp from the real fps/pix_fmt
        # instead of falling back to defaults (wrong seek → frame desync).
        self.fps = float(fps or 30.0)
        self.pix_fmt = pix_fmt
        seek_time = max(0, int(start_frame)) / self.fps
        # Truncate to milliseconds (never round up): ffmpeg's accurate seek
        # returns the first frame with ts >= seek_time, so a rounded-up value
        # such as 36.6667 lands one frame PAST the target (36.66667 -> frame
        # 1101 instead of 1100).  A 1ms truncation is always below the
        # target frame ts and at least ~32ms above the previous frame's.
        seek_time = math.floor(seek_time * 1000) / 1000
        self._process = subprocess.Popen(
            [ffmpeg, "-v", "error",
             "-ss", f"{seek_time:.3f}", "-i", str(self.path),
             "-f", "rawvideo", "-pix_fmt", pix_fmt, "pipe:1"],
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        )

    def read_codes(self):
        """Read one stored little-endian uint16 frame without colorization."""
        import numpy as np

        if self._process.stdout is None:
            return None
        chunks: list[bytes] = []
        remaining = self._frame_bytes
        while remaining:
            data = self._process.stdout.read(remaining)
            if not data:
                return None
            chunks.append(data)
            remaining -= len(data)
        return np.frombuffer(b"".join(chunks), dtype="<u2").reshape(
            self.height, self.width,
        ).copy()

    def close(self) -> None:
        process = self._process
        if process.poll() is None:
            # Same blocked-pipe situation as DepthVideoReader.close(); a
            # graceful SIGTERM flush hangs for seconds on a full stdout pipe.
            process.terminate()
            try:
                process.wait(timeout=0.3)
            except subprocess.TimeoutExpired:
                process.kill()
                try:
                    process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    pass
        if process.stdout:
            process.stdout.close()
        if process.stderr:
            process.stderr.close()


def _to_mp4(source: Path, destination: Path) -> Path:
    """Put a video in the v2.1 MP4 path, preserving codec when possible."""
    destination.parent.mkdir(parents=True, exist_ok=True)
    if source.resolve() == destination.resolve():
        return source
    if source.suffix.lower() == ".mp4":
        shutil.copy2(source, destination)
        return destination
    try:
        completed = subprocess.run(
            ["ffmpeg", "-y", "-loglevel", "error", "-i", str(source),
             "-map", "0:v:0", "-c", "copy", "-movflags", "+faststart",
             str(destination)],
            capture_output=True, timeout=1800, check=False,
        )
        if completed.returncode == 0 and destination.is_file():
            return destination
    except (OSError, subprocess.TimeoutExpired):
        pass
    # Keep the original codec/container rather than discarding a depth stream.
    fallback = destination.with_suffix(source.suffix.lower())
    shutil.copy2(source, fallback)
    return fallback


def _canonicalize_depth_video(source: Path, destination: Path,
                              fps: float) -> dict[str, Any]:
    """Transcode a legacy gray depth video to canonical log depth codes."""
    import numpy as np

    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg is required to canonicalize depth video")
    destination = Path(destination)
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_name(
        f".{destination.stem}.tmp-{uuid4().hex}{destination.suffix or '.mp4'}"
    )
    reader = DepthVideoReader(source)
    process = None
    written = 0
    try:
        command = [
            ffmpeg, "-y", "-loglevel", "error",
            "-f", "rawvideo", "-pix_fmt", "gray12le",
            "-s", f"{reader.width}x{reader.height}",
            "-r", str(float(fps or reader.fps or 30.0)), "-i", "pipe:0",
            "-an", "-c:v", "libx265", *depth_video_encoder_args(),
            "-movflags", "+faststart", str(temporary),
        ]
        process = subprocess.Popen(
            command, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
        )
        assert process.stdin is not None
        while True:
            codes = reader.read_canonical_codes()
            if codes is None:
                break
            process.stdin.write(np.ascontiguousarray(codes, dtype="<u2").tobytes())
            written += 1
        process.stdin.close()
        stderr = process.stderr.read() if process.stderr is not None else b""
        return_code = process.wait(timeout=1800)
        if return_code != 0 or written <= 0 or not temporary.is_file():
            detail = (stderr.decode(errors="replace")
                      or "ffmpeg returned no output").strip()
            raise RuntimeError(f"Depth video canonicalization failed: {detail[-1000:]}")
        probe = _probe_video(temporary, float(fps or reader.fps or 30.0))
        if probe.get("codec") not in {"hevc", "h265"} \
                or probe.get("pix_fmt") != "gray12le":
            raise RuntimeError(
                "Depth video canonicalization verification failed: "
                f"codec={probe.get('codec')!r}, pix_fmt={probe.get('pix_fmt')!r}"
            )
        temporary.replace(destination)
        probe.update({
            "frames": int(probe.get("frames") or written),
            "fps": float(probe.get("fps") or fps or reader.fps or 30.0),
            "depth_encoding": DEPTH_VIDEO_ENCODING,
            "depth_min_mm": DEPTH_MIN_MM,
            "depth_max_mm": DEPTH_MAX_MM,
            "depth_qmax": DEPTH_QMAX,
            "depth_qp": DEPTH_QP,
        })
        return probe
    except (BrokenPipeError, OSError, subprocess.TimeoutExpired) as exc:
        if process is not None and process.poll() is None:
            process.kill()
            process.wait()
        raise RuntimeError(f"Depth video canonicalization failed: {exc}") from exc
    finally:
        reader.close()
        if process is not None and process.poll() is None:
            process.kill()
            process.wait()
        temporary.unlink(missing_ok=True)


def _source_feature(info: dict, source: str, probe: dict, metric_depth: bool) -> dict:
    key = f"observation.images.{source}"
    features = info.get("features") if isinstance(info.get("features"), dict) else {}
    original = features.get(key) if isinstance(features.get(key), dict) else {}
    original_video = original.get("video_info")
    video_info = dict(original_video) if isinstance(original_video, dict) else {}
    is_depth = bool(video_info.get("video.is_depth_map")) or metric_depth
    video_info.update({
        "video.fps": probe["fps"],
        "video.height": probe["height"],
        "video.width": probe["width"],
        "video.channels": 1 if is_depth else 3,
        "video.codec": probe.get("codec") or video_info.get("video.codec") or "",
        "video.pix_fmt": probe.get("pix_fmt") or video_info.get("video.pix_fmt") or "",
        "video.is_depth_map": is_depth,
        "video.is_depth_visualization": bool(is_depth is False and is_depth_source(source)),
        "has_audio": False,
    })
    if metric_depth:
        video_info.update({
            "video.depth_min_mm": DEPTH_MIN_MM,
            "video.depth_max_mm": DEPTH_MAX_MM,
            "video.depth_qmax": DEPTH_QMAX,
            "video.depth_qp": DEPTH_QP,
            "video.depth_quantization": "log",
            "video.depth_encoding": probe.get("depth_encoding") or DEPTH_VIDEO_ENCODING,
        })
    feature = dict(original)
    feature.update({
        "dtype": "video",
        "shape": [probe["height"], probe["width"], video_info["video.channels"]],
        "names": ["height", "width", "channel"],
        "video_info": video_info,
    })
    return feature


def _read_task_rows(meta_root: Path, fallback: str) -> tuple[dict[int, str], str]:
    tasks: dict[int, str] = {}
    json_path = meta_root / "tasks.json"
    json_value = _read_json(json_path)
    if isinstance(json_value, dict):
        json_value = json_value.get("tasks", json_value)
    if isinstance(json_value, list):
        for default_index, row in enumerate(json_value):
            if isinstance(row, str):
                index, text = default_index, row.strip()
            elif isinstance(row, dict):
                try:
                    index = int(row.get("task_index", row.get("task_id", default_index)))
                except (TypeError, ValueError):
                    index = default_index
                text = str(row.get("task") or row.get("description") or "").strip()
            else:
                continue
            if text:
                tasks[index] = text
    for filename in ("tasks.jsonl",):
        if tasks:
            break
        path = meta_root / filename
        if not path.is_file():
            continue
        for line in path.read_text(encoding="utf-8").splitlines():
            try:
                row = json.loads(line)
            except (ValueError, TypeError):
                continue
            if not isinstance(row, dict):
                continue
            try:
                index = int(row.get("task_index", row.get("task_id", 0)))
            except (TypeError, ValueError):
                index = 0
            text = str(row.get("task") or row.get("description") or "").strip()
            if text:
                tasks[index] = text
    if not tasks:
        tasks[0] = fallback or "default_recording"
    return tasks, tasks[min(tasks)]


def _episode_stats(frame_table) -> dict[str, dict[str, Any]]:
    stats: dict[str, dict[str, Any]] = {}
    for column in frame_table.columns:
        try:
            series = frame_table[column]
            if not getattr(series.dtype, "kind", "") in "biuf":
                continue
            values = series.to_numpy(dtype=float)
            if not len(values):
                continue
            import numpy as np
            finite = values[np.isfinite(values)]
            if not len(finite):
                continue
            stats[column] = {
                "min": float(finite.min()), "max": float(finite.max()),
                "mean": float(finite.mean()), "std": float(finite.std()),
                "count": [int(len(finite))],
            }
        except Exception:
            continue
    return stats


def _write_episode_index(meta_root: Path,
                         rows: Iterable[dict[str, Any]]) -> None:
    """Write the canonical sharded Parquet episode index for one upload."""
    import pyarrow as pa
    import pyarrow.parquet as pq

    episodes_root = Path(meta_root) / "episodes"
    if episodes_root.is_dir():
        shutil.rmtree(episodes_root)
    normalized = [dict(row) for row in rows]
    normalized.sort(key=lambda row: int(row.get("episode_index", 10**9)))
    for row in normalized:
        episode_index = int(row.get("episode_index", 0))
        chunk = episode_index // EPISODES_METADATA_CHUNK_SIZE
        target = (episodes_root / f"chunk-{chunk:03d}"
                  / f"episode_{episode_index:06d}.parquet")
        target.parent.mkdir(parents=True, exist_ok=True)
        pq.write_table(
            pa.Table.from_pylist([row]),
            target,
        )


def normalize_extracted_dataset(root: Path, batch_name: str) -> dict[str, Any]:
    """Normalize one extracted collector archive to LeRobot v2.1 in-place."""
    root = Path(root)
    meta_root = root / "meta"
    meta_root.mkdir(parents=True, exist_ok=True)
    info = normalize_metadata_sources(_read_json(meta_root / "info.json"))
    fallback_task = str(info.get("task_name") or batch_name)
    data_paths = _data_parquets(root)
    try:
        frame_table = _load_frame_table(data_paths)
    except RuntimeError:
        frame_table = _legacy_frame_table(root, info, batch_name)
    episode_index = _infer_episode_index(root, batch_name, frame_table)
    frame_table = frame_table.copy()
    frame_table["episode_index"] = episode_index
    if "frame_index" in frame_table.columns:
        frame_table = frame_table.sort_values("frame_index", kind="stable")
    frame_count = int(len(frame_table))
    fps = float(info.get("fps") or 30.0)

    # Replace all source data parquet files with one v2.1 episode parquet.
    data_target = root / "data" / "chunk-000" / f"episode_{episode_index:06d}.parquet"
    data_target.parent.mkdir(parents=True, exist_ok=True)
    frame_table.to_parquet(data_target, index=False)
    for path in data_paths:
        if path.resolve() != data_target.resolve():
            path.unlink(missing_ok=True)
    for directory in sorted((root / "data").glob("**/*"), reverse=True):
        if directory.is_dir() and directory != data_target.parent:
            try:
                directory.rmdir()
            except OSError:
                pass

    # Normalize one video per source to the v2.1 path template.
    source_videos = iter_video_streams(root / "videos")
    normalized_video_sources = list(source_videos)
    grouped: dict[str, list[Path]] = {}
    for source, path in source_videos:
        grouped.setdefault(source, []).append(path)
    video_features: dict[str, dict] = {}
    video_extensions: dict[str, str] = {}
    video_meta: dict[str, dict[str, Any]] = {}
    videos_root = root / "videos"
    for source, paths in sorted(grouped.items()):
        if len(paths) > 1:
            raise RuntimeError(
                f"Multiple video files found for source {source}; "
                "one uploaded batch must contain one file per source"
            )
        source_path = paths[0]
        probe = _probe_video(source_path, fps)
        # A pure depth stream is metric only when its metadata/pixel format
        # proves it.  A normal yuv420p/colourized depth preview is not 3D data.
        old_feature = (info.get("features") or {}).get(
            f"observation.images.{source}", {})
        old_video = old_feature.get("video_info") if isinstance(old_feature, dict) else {}
        metric_depth = bool(isinstance(old_video, dict)
                            and old_video.get("video.is_depth_map"))
        metric_depth = metric_depth or (
            is_depth_source(source)
            and probe.get("pix_fmt") in {"gray12le", "gray16le", "gray10le"}
        )
        destination = (videos_root / f"observation.images.{_safe_source_key(source)}"
                       / "chunk-000" / f"episode_{episode_index:06d}.mp4")
        if metric_depth and _depth_video_mode(old_video) != "log":
            # Legacy gray16/gray12 streams may carry direct millimetres.  Do
            # not relabel them as log codes without converting the samples.
            probe = _canonicalize_depth_video(source_path, destination, fps)
            final_path = destination
        else:
            final_path = _to_mp4(source_path, destination)
        video_extensions[source] = final_path.suffix.lstrip(".")
        video_meta[source] = {
            "frames": probe["frames"] or frame_count,
            "fps": probe["fps"],
            "metric_depth": metric_depth,
            "source": source,
        }
        video_features[f"observation.images.{source}"] = _source_feature(
            info, source, probe, metric_depth,
        )

    # Legacy collectors may upload metric depth as depth/<source>/*.png.
    # Convert it before the legacy tree is removed.  New projects therefore
    # keep only one independent depth video beside RGB streams, while the
    # resolver continues to read old meta/depth PNGs from already-migrated
    # projects.
    depth_root = root / "depth"
    depth_directories = depth_png_directories(depth_root)
    for source, depth_dir, frames in depth_directories:
        if source in video_meta:
            # An explicitly supplied depth video is authoritative.  The PNG
            # copy is redundant and is removed only after the whole batch has
            # passed the output verification below.
            continue
        destination = (videos_root / f"observation.images.{_safe_source_key(source)}"
                       / "chunk-000" / f"episode_{episode_index:06d}.mp4")
        depth_probe = encode_depth_png_sequence(
            depth_dir, frames, destination, fps=float(info.get("fps") or fps or 30.0),
        )
        depth_scale = 0.001
        depth_feature = _source_feature(info, source, depth_probe, True)
        depth_video_info = depth_feature.setdefault("video_info", {})
        depth_video_info.update({
            "video.channels": 1,
            "video.is_depth_map": True,
            "video.is_depth_visualization": False,
            "video.depth_scale": depth_scale,
            "video.depth_encoding": DEPTH_VIDEO_ENCODING,
            "video.depth_min_mm": DEPTH_MIN_MM,
            "video.depth_max_mm": DEPTH_MAX_MM,
            "video.depth_qmax": DEPTH_QMAX,
            "video.depth_qp": DEPTH_QP,
            "video.depth_quantization": "log",
            "has_audio": False,
        })
        depth_feature["shape"] = [
            int(depth_probe.get("height") or 0),
            int(depth_probe.get("width") or 0), 1,
        ]
        video_features[f"observation.images.{source}"] = depth_feature
        video_extensions[source] = "mp4"
        video_meta[source] = {
            "frames": int(depth_probe.get("frames") or len(frames)),
            "fps": float(depth_probe.get("fps") or fps or 30.0),
            "metric_depth": True,
            "source": source,
            "depth_scale": depth_scale,
        }
        normalized_video_sources.append((source, destination))

    # Do not retain raw depth PNGs in a newly normalized batch.  This is
    # reached only after every sequence has produced and verified its video;
    # if conversion fails, normalization aborts and the upload staging tree
    # remains available for retry/diagnosis.
    if depth_root.is_dir():
        shutil.rmtree(depth_root)

    # Remove any source/legacy video tree only after all files have been
    # copied.  The canonical v2.1 layout is videos/<video_key>/chunk-000/.
    for child in list(videos_root.iterdir()) if videos_root.is_dir() else []:
        if not child.name.startswith("observation.images."):
            if child.is_dir():
                shutil.rmtree(child, ignore_errors=True)
            else:
                child.unlink(missing_ok=True)
            continue
        for nested in list(child.iterdir()) if child.is_dir() else []:
            if nested.name != "chunk-000":
                if nested.is_dir():
                    shutil.rmtree(nested, ignore_errors=True)
                else:
                    nested.unlink(missing_ok=True)

    tasks, default_task = _read_task_rows(meta_root, fallback_task)
    # Keep an explicit task file, but remove v3-only pooled episode metadata.
    for stale in (meta_root / "tasks.parquet",):
        stale.unlink(missing_ok=True)
    episode_row: dict[str, Any] = {
        "episode_index": episode_index,
        "tasks": [default_task],
        "task_index": min(tasks),
        "length": frame_count,
        "data/chunk_index": 0,
        "data/file_index": episode_index,
        "dataset_from_index": 0,
        "dataset_to_index": frame_count,
    }
    for source, values in sorted(video_meta.items()):
        key = f"observation.images.{source}"
        count = int(values.get("frames") or frame_count)
        stream_fps = float(values.get("fps") or fps)
        episode_row[f"videos/{key}/chunk_index"] = 0
        episode_row[f"videos/{key}/file_index"] = episode_index
        episode_row[f"videos/{key}/from_timestamp"] = 0.0
        episode_row[f"videos/{key}/to_timestamp"] = count / stream_fps if stream_fps else 0.0
    episode_stats = _episode_stats(frame_table)
    episode_row["stats"] = episode_stats
    _write_episode_index(meta_root, [episode_row])
    (meta_root / "episodes_stats.jsonl").unlink(missing_ok=True)
    tasks_path = meta_root / "tasks.json"
    tasks_path.write_text(
        json.dumps([{"task_index": index, "task": text}
                    for index, text in sorted(tasks.items())],
                   ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (meta_root / "tasks.jsonl").unlink(missing_ok=True)

    # Move collector-only metadata under meta/, keeping the dataset root at
    # exactly data/meta/videos.  Calibration is metadata, not a fourth root.
    calibration_root = root / "calibration"
    if calibration_root.is_dir():
        target = meta_root / "calibration"
        target.mkdir(parents=True, exist_ok=True)
        for item in calibration_root.iterdir():
            shutil.move(str(item), str(target / item.name))
        calibration_root.rmdir()
    for filename in ("metadata.json", "timestamps.json"):
        source = root / filename
        if source.is_file():
            target = meta_root / "collector" / filename
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(str(source), str(target))

    # Keep the dataset root limited to data/meta/videos without hiding raw
    # depth frames from the processing pipeline.  PNG depth is auxiliary
    # source material, so it lives below meta/depth while the synchronized
    # video stream remains under videos/observation.images.<source>/.
    for item in list(root.iterdir()):
        if item.name in {"data", "meta", "videos"}:
            continue
        if item.name == "depth":
            target = meta_root / "depth"
        else:
            target = meta_root / "collector" / item.name
        target.parent.mkdir(parents=True, exist_ok=True)
        if target.exists():
            raise RuntimeError(f"Metadata destination already exists: {target}")
        shutil.move(str(item), str(target))

    # Collector metadata is part of the dataset contract.  Normalize it too,
    # otherwise a historical upload could reintroduce old source/device names
    # after the active videos were already written with canonical keys.
    for metadata_path in meta_root.rglob("*.json"):
        metadata = _read_json(metadata_path)
        if not metadata:
            continue
        normalized_metadata = normalize_metadata_sources(metadata)
        _write_json(metadata_path, normalized_metadata)
    info = normalize_metadata_sources(info)

    features = dict(info.get("features") or {}) if isinstance(info.get("features"), dict) else {}
    features.update(video_features)
    normalized = dict(info)
    normalized.update({
        "format": "lerobot_v2.1",
        "codebase_version": "v2.1",
        "total_episodes": 1,
        "total_frames": frame_count,
        "total_tasks": len(tasks),
        "chunks_size": 1000,
        "fps": fps,
        "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet",
        "video_path": "videos/{video_key}/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.{ext}",
        "features": features,
        "video_extensions": video_extensions,
        "extensions": {
            **(info.get("extensions") or {}),
            "episodes_file": "meta/episodes/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet",
            "stats_file": "meta/stats.json",
            "tasks_file": "meta/tasks.json",
        },
    })
    _write_json(meta_root / "info.json", normalized)

    # A standalone upload_manifest is useful for diagnostics but belongs in
    # meta/ so it does not create a fourth dataset-root directory.
    return {
        "episode_index": episode_index,
        "frame_count": frame_count,
        "fps": fps,
        "video_sources": sorted(source for source, _ in normalized_video_sources),
        "camera_sources": sorted(source for source, _ in normalized_video_sources
                                  if not is_depth_source(source)),
        "depth_sources": sorted(source for source, _ in normalized_video_sources
                                 if is_depth_source(source)),
        "video_meta": video_meta,
        "info": normalized,
    }