tmp/video-velocity-model/data_processing/traj_to_mp4.py

#!/usr/bin/env python3
"""
BridgeData v2 轨迹 → MP4 视频还原

将 scripted_raw 目录中的每条机器人轨迹（图像序列 + 时间戳）编码为 MP4 视频。
支持基于真实时间戳的 VFR → CFR 转换，多相机任务，并行编码。

Usage:
    # 预览模式
    python data_processing/traj_to_mp4.py --dry-run

    # 编码单个任务（测试用）
    python data_processing/traj_to_mp4.py \\
        --task-filter "*rigid_objects*" --workers 1 \\
        --output-dir /path/to/output

    # 全量编码
    python data_processing/traj_to_mp4.py \\
        --output-dir /path/to/output --workers 8
"""

from __future__ import annotations

import argparse
import fnmatch
import logging
import os
import pickle
import re
import subprocess
import sys
import tempfile
import time
from collections.abc import Iterator
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Optional

import numpy as np

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

DEFAULT_INPUT_DIR = "/home/sz128/scratch_sz128/datasets/bridgedatav2/raw"
DEFAULT_FFMPEG_BIN = "/home/sz128/ffmpeg-n7.1-latest-linux64-gpl-7.1/bin/ffmpeg"
DEFAULT_CRF = 23
DEFAULT_WORKERS = 8
DEFAULT_FFMPEG_THREADS = 2

IMAGE_DIR_PATTERN = re.compile(r"^images(\d+)$")
DEPTH_DIR_PATTERN = re.compile(r"^depth_images(\d+)$")

logger = logging.getLogger("traj_to_mp4")


# ---------------------------------------------------------------------------
# Discovery
# ---------------------------------------------------------------------------


def discover_trajectories(input_dir: Path) -> list[Path]:
    """Find all trajectory directories under the input root.

    Recursively locates every ``raw/traj_group0/traj{N}/`` directory,
    regardless of how many intermediate directory levels exist above it.

    Uses ``os.walk`` with pruning to skip image directories for speed.

    Works with both layouts::

        # Flat (scripted_raw)
        {input_dir}/{task}/{episode}/raw/traj_group0/traj{N}/

        # Deep (raw)
        {input_dir}/{dataset}/{collection}/.../{episode}/raw/traj_group0/traj{N}/

    Returns a list of absolute paths to each ``traj{N}`` directory.
    """
    traj_dirs: list[Path] = []
    for root, dirs, _files in os.walk(input_dir):
        # Prune image directories — they are leaf dirs that never contain
        # traj_group0, and skipping them avoids scanning millions of files.
        dirs[:] = [
            d
            for d in dirs
            if not (d.startswith("images") or d.startswith("depth_images"))
        ]
        if os.path.basename(root) == "traj_group0" and os.path.basename(
            os.path.dirname(root)
        ) == "raw":
            for traj_name in sorted(dirs):
                if traj_name.startswith("traj"):
                    traj_dirs.append((Path(root) / traj_name).resolve())
    return traj_dirs


def filter_trajectories(
    traj_dirs: list[Path],
    task_filter: Optional[str] = None,
    input_dir: Optional[Path] = None,
) -> list[Path]:
    """Filter trajectory list by optional glob against the relative path.

    The filter is matched against the trajectory's path relative to
    ``input_dir``, so any path segment can be targeted (e.g.
    ``*toykitchen1*``, ``*bridge_data_v2*``, ``*sweep_12-03*``).

    If *task_filter* contains no glob characters (``*?[]``), it is
    automatically wrapped with ``*...*`` for substring matching.
    """
    if not task_filter:
        return traj_dirs

    has_glob = any(c in task_filter for c in "*?[]")

    filtered: list[Path] = []
    for p in traj_dirs:
        if input_dir is not None:
            try:
                rel = str(p.relative_to(input_dir))
            except ValueError:
                rel = str(p)
        else:
            rel = str(p)

        pattern = task_filter if has_glob else f"*{task_filter}*"
        if fnmatch.fnmatch(rel, pattern):
            filtered.append(p)
    return filtered


# ---------------------------------------------------------------------------
# Camera / image helpers
# ---------------------------------------------------------------------------


def list_camera_dirs(traj_dir: Path, include_depth: bool = False) -> list[str]:
    """Return sorted list of camera directory names (e.g. ``images0``)."""
    cameras: list[str] = []
    for entry in sorted(traj_dir.iterdir()):
        if not entry.is_dir():
            continue
        if IMAGE_DIR_PATTERN.match(entry.name):
            cameras.append(entry.name)
        elif include_depth and DEPTH_DIR_PATTERN.match(entry.name):
            cameras.append(entry.name)
    return cameras


def collect_image_paths(traj_dir: Path, camera_dir: str) -> list[Path]:
    """Return sorted list of absolute image paths from a camera directory.

    Images are sorted by the numeric index extracted from ``im_{N}.{ext}``.
    """
    cam_path = traj_dir / camera_dir
    if not cam_path.is_dir():
        return []

    images: list[tuple[int, Path]] = []
    for f in cam_path.iterdir():
        if not f.is_file():
            continue
        m = re.match(r"im_(\d+)\.", f.name)
        if m:
            images.append((int(m.group(1)), f.resolve()))
    images.sort(key=lambda x: x[0])
    return [p for _, p in images]


# ---------------------------------------------------------------------------
# Timestamp extraction
# ---------------------------------------------------------------------------


def extract_timestamps(traj_dir: Path) -> np.ndarray:
    """Load ``obs_dict.pkl`` and return the ``time_stamp`` array (seconds)."""
    pkl_path = traj_dir / "obs_dict.pkl"
    if not pkl_path.is_file():
        raise FileNotFoundError(f"obs_dict.pkl not found in {traj_dir}")
    with open(pkl_path, "rb") as f:
        obs = pickle.load(f)
    ts = np.asarray(obs["time_stamp"], dtype=np.float64)
    return ts


# ---------------------------------------------------------------------------
# FFmpeg concat file
# ---------------------------------------------------------------------------


def write_concat_file(
    image_paths: list[Path],
    timestamps: np.ndarray,
    concat_path: Path,
) -> float:
    """Write an FFmpeg concat demuxer file with per-frame durations.

    Returns the mean frame interval (seconds), which can be used as a
    fallback FPS suggestion.
    """
    n_images = len(image_paths)
    n_ts = len(timestamps)

    if n_images != n_ts:
        logger.warning(
            "Image count (%d) != timestamp count (%d), using min", n_images, n_ts
        )
        n_frames = min(n_images, n_ts)
    else:
        n_frames = n_images

    # Compute frame durations from timestamps
    if n_frames >= 2:
        durations = np.diff(timestamps[:n_frames])
        # Replace negative or zero durations with a small default
        durations = np.clip(durations, 0.001, None)
        mean_interval = float(np.mean(durations))
    else:
        durations = np.array([])
        mean_interval = 0.2  # fallback

    lines = ["ffconcat version 1.0\n"]
    for i in range(n_frames):
        lines.append(f"file '{image_paths[i]}'\n")
        if i < n_frames - 1:
            lines.append(f"duration {durations[i]:.6f}\n")

    concat_path.write_text("".join(lines), encoding="utf-8")
    return mean_interval


# ---------------------------------------------------------------------------
# FFmpeg encoding
# ---------------------------------------------------------------------------


def build_ffmpeg_cmd(
    concat_file: Path,
    output_path: Path,
    *,
    target_fps: Optional[float],
    crf: int,
    ffmpeg_bin: str,
    ffmpeg_threads: int,
) -> list[str]:
    """Build the FFmpeg command line for concat → MP4 encoding."""
    cmd = [
        ffmpeg_bin,
        "-y",
        "-f",
        "concat",
        "-safe",
        "0",
        "-i",
        str(concat_file),
        "-c:v",
        "libx264",
        "-crf",
        str(crf),
        "-preset",
        "fast",
        "-pix_fmt",
        "yuv420p",
        "-movflags",
        "+faststart",
        "-threads",
        str(ffmpeg_threads),
    ]

    if target_fps is not None and target_fps > 0:
        cmd.extend(["-vf", f"fps={target_fps}"])

    cmd.append(str(output_path))
    return cmd


def encode_traj(
    traj_dir: Path,
    output_dir: Path,
    *,
    target_fps: Optional[float],
    crf: int,
    ffmpeg_bin: str,
    ffmpeg_threads: int,
    include_depth: bool = False,
    skip_existing: bool = False,
) -> tuple[str, str]:
    """Encode one trajectory to MP4.

    Returns ``(status, message)`` where status is ``"ok"``, ``"skipped"``,
    or ``"error"``.
    """
    traj_name = traj_dir.name
    # Flat output: all videos go directly under output_dir
    # Filename: {task}_{episode}_{traj_name}[_{camera}].mp4
    episode_dir = traj_dir.parent.parent.parent  # traj_group0/raw/episode
    task_dir = episode_dir.parent
    prefix = f"{task_dir.name}_{episode_dir.name}_{traj_name}"
    out_parent = output_dir
    out_parent.mkdir(parents=True, exist_ok=True)

    # Load timestamps once
    try:
        timestamps = extract_timestamps(traj_dir)
    except Exception as e:
        return ("error", f"{traj_dir}: failed to read timestamps: {e}")

    cameras = list_camera_dirs(traj_dir, include_depth=include_depth)
    if not cameras:
        return ("error", f"{traj_dir}: no image directories found")

    for cam in cameras:
        # Determine output filename
        if len(cameras) == 1 and cam == "images0":
            out_name = f"{prefix}.mp4"
        else:
            out_name = f"{prefix}_{cam}.mp4"
        out_path = out_parent / out_name

        if skip_existing and out_path.is_file():
            continue

        # Collect images
        image_paths = collect_image_paths(traj_dir, cam)
        if not image_paths:
            return ("error", f"{traj_dir}/{cam}: no images found")

        # Write concat file
        concat_fd, concat_path = tempfile.mkstemp(
            suffix=".txt", prefix="ffconcat_"
        )
        try:
            mean_interval = write_concat_file(image_paths, timestamps, Path(concat_path))

            # Determine FPS
            fps = target_fps
            if fps is None:
                # Auto: use rounded mean interval to get ~4-5 FPS
                fps = round(1.0 / mean_interval) if mean_interval > 0 else 4
                fps = max(1, min(fps, 60))

            # Build and run FFmpeg
            cmd = build_ffmpeg_cmd(
                Path(concat_path),
                out_path,
                target_fps=fps,
                crf=crf,
                ffmpeg_bin=ffmpeg_bin,
                ffmpeg_threads=ffmpeg_threads,
            )

            result = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                timeout=300,
                check=False,
            )

            if result.returncode != 0:
                stderr_tail = result.stderr.strip().split("\n")[-5:]
                return (
                    "error",
                    f"{traj_dir}/{cam}: ffmpeg failed:\n" + "\n".join(stderr_tail),
                )

            if not out_path.is_file():
                return ("error", f"{traj_dir}/{cam}: output file missing after ffmpeg")

        finally:
            os.close(concat_fd)
            try:
                os.unlink(concat_path)
            except OSError:
                pass

    return ("ok", f"{traj_dir}: {len(cameras)} camera(s) encoded")


# ---------------------------------------------------------------------------
# Worker entry point (module-level for ProcessPoolExecutor)
# ---------------------------------------------------------------------------


def _encode_worker(args: tuple) -> tuple[str, str]:
    """Worker function (must be picklable)."""
    (
        traj_dir,
        output_dir,
        target_fps,
        crf,
        ffmpeg_bin,
        ffmpeg_threads,
        include_depth,
        skip_existing,
    ) = args
    return encode_traj(
        Path(traj_dir),
        Path(output_dir),
        target_fps=target_fps,
        crf=crf,
        ffmpeg_bin=ffmpeg_bin,
        ffmpeg_threads=ffmpeg_threads,
        include_depth=include_depth,
        skip_existing=skip_existing,
    )


# ---------------------------------------------------------------------------
# Main orchestrator
# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(
        description="BridgeData v2 轨迹 → MP4 视频还原",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--input-dir",
        default=DEFAULT_INPUT_DIR,
        help=f"数据根目录 (default: {DEFAULT_INPUT_DIR})",
    )
    parser.add_argument(
        "--output-dir",
        default="/home/sz128/scratch_sz128/datasets/bridgedatav2_video/data/raw",
    )
    parser.add_argument(
        "--fps",
        default="auto",
        help="目标 FPS: 'auto' (自动从时间戳计算), 'vfr' (保留可变帧率), 或具体数值如 5",
    )
    parser.add_argument(
        "--crf",
        type=int,
        default=DEFAULT_CRF,
        help=f"H.264 CRF 值, 越小质量越高 (default: {DEFAULT_CRF})",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=DEFAULT_WORKERS,
        help=f"并行 worker 数 (default: {DEFAULT_WORKERS})",
    )
    parser.add_argument(
        "--ffmpeg-bin",
        default=DEFAULT_FFMPEG_BIN,
        help=f"FFmpeg 可执行文件路径 (default: {DEFAULT_FFMPEG_BIN})",
    )
    parser.add_argument(
        "--ffmpeg-threads",
        type=int,
        default=DEFAULT_FFMPEG_THREADS,
        help=f"每个 FFmpeg 进程的线程数 (default: {DEFAULT_FFMPEG_THREADS})",
    )
    parser.add_argument(
        "--skip-existing",
        action="store_true",
        help="跳过已存在的输出文件",
    )
    parser.add_argument(
        "--include-depth",
        action="store_true",
        help="同时编码深度图（depth_images*）为视频",
    )
    parser.add_argument(
        "--task-filter",
        default=None,
        help="只处理匹配的任务目录名（glob pattern），如 '*rigid_objects*'",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="只统计不编码",
    )
    args = parser.parse_args()

    # Parse FPS
    if args.fps == "auto":
        target_fps: Optional[float] = None
    elif args.fps.lower() == "vfr":
        target_fps = None  # Will be handled as "no -vf fps="
        # Actually VFR means don't convert: we'll use a sentinel
        target_fps = -1.0  # negative = skip fps filter
    else:
        try:
            target_fps = float(args.fps)
            if target_fps <= 0:
                raise ValueError
        except ValueError:
            parser.error(f"Invalid --fps value: {args.fps}")

    # Resolve paths
    input_dir = Path(args.input_dir).expanduser().resolve()
    output_dir = Path(args.output_dir).expanduser().resolve()

    if not input_dir.is_dir():
        print(f"ERROR: input directory not found: {input_dir}", file=sys.stderr)
        sys.exit(1)

    # Verify ffmpeg
    ffmpeg_bin = args.ffmpeg_bin
    try:
        subprocess.run(
            [ffmpeg_bin, "-version"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=True,
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        print(f"ERROR: ffmpeg not found at {ffmpeg_bin}", file=sys.stderr)
        sys.exit(1)

    # Discovery
    print("Discovering trajectories...")
    all_trajs = discover_trajectories(input_dir)
    trajs = filter_trajectories(all_trajs, args.task_filter, input_dir=input_dir)
    print(f"  Found {len(all_trajs)} total, {len(trajs)} after filtering")

    if not trajs:
        print("No trajectories to process.")
        sys.exit(0)

    # Count cameras and images for estimation
    total_cameras = 0
    total_frames = 0
    for t in trajs[:10]:  # sample first 10 for quick estimate
        cams = list_camera_dirs(t, include_depth=args.include_depth)
        if cams:
            imgs = collect_image_paths(t, cams[0])
            total_cameras += len(cams)
            total_frames += len(imgs) * len(cams)
    avg_cameras = max(1, total_cameras / min(10, len(trajs)))
    avg_frames_per_cam = total_frames / max(1, total_cameras)
    est_videos = int(len(trajs) * avg_cameras)
    print(f"  Estimated videos: ~{est_videos}")
    print(f"  Estimated frames/video: ~{avg_frames_per_cam:.0f}")
    print(f"  Output dir: {output_dir}")

    if args.dry_run:
        print("\n[Dry run] Encoding would start now. Remove --dry-run to execute.")
        sys.exit(0)

    # Prepare worker arguments
    worker_args = [
        (
            str(t),
            str(output_dir),
            target_fps,
            args.crf,
            ffmpeg_bin,
            args.ffmpeg_threads,
            args.include_depth,
            args.skip_existing,
        )
        for t in trajs
    ]

    # Progress callback
    success_count = 0
    skip_count = 0
    fail_count = 0
    failed_items: list[str] = []

    started_at = time.perf_counter()
    print(f"\nEncoding with {args.workers} workers...")

    with ProcessPoolExecutor(max_workers=args.workers) as pool:
        futures = {
            pool.submit(_encode_worker, wa): wa[0] for wa in worker_args
        }

        for i, future in enumerate(as_completed(futures), 1):
            traj_path = futures[future]
            try:
                status, message = future.result()
            except Exception as e:
                status, message = "error", f"{traj_path}: worker exception: {e}"

            if status == "ok":
                success_count += 1
            elif status == "skipped":
                skip_count += 1
            else:
                fail_count += 1
                failed_items.append(message)

            # Progress
            pct = i / len(futures) * 100
            elapsed = time.perf_counter() - started_at
            rate = i / max(1, elapsed)
            eta = (len(futures) - i) / max(0.001, rate)
            print(
                f"\r[{i}/{len(futures)} {pct:.0f}%] "
                f"ok={success_count} skip={skip_count} fail={fail_count} "
                f"| {rate:.1f}/s ETA {eta:.0f}s   ",
                end="",
                flush=True,
            )

    elapsed = time.perf_counter() - started_at
    print(f"\n\nDone in {elapsed:.0f}s.")

    # Write failure log
    if failed_items:
        fail_log = output_dir / "failed_trajs.log"
        fail_log.write_text("\n".join(failed_items), encoding="utf-8")
        print(f"Failures written to {fail_log}")

    print(
        f"Summary: {success_count} ok, {skip_count} skipped, {fail_count} failed "
        f"(total {len(trajs)})"
    )


if __name__ == "__main__":
    main()