scripts/prepare_ctrlworld_single_arm_multiview.py

"""
Prepare single_arm/multiview LeRoBot dataset into Ctrl-World annotation format.

For each episode:
  - Reads parquet (states, joints, gripper, instruction)
  - Reads original 15fps videos and downsamples to 5fps (every 3rd frame)
  - Resizes 180 -> 192 height (Ctrl-World expects 192x320)
  - Saves annotation JSON + downsampled videos in Ctrl-World expected layout

Output structure (per subset):
  {output_dir}/{subset}/annotation/val/{episode_id}.json
  {output_dir}/{subset}/videos/val/{episode_id}/0.mp4  (exterior_1_left)
  {output_dir}/{subset}/videos/val/{episode_id}/1.mp4  (exterior_2_left)
  {output_dir}/{subset}/videos/val/{episode_id}/2.mp4  (wrist_left)
"""

import argparse
import json
import os
import subprocess
from pathlib import Path

import numpy as np
import pandas as pd


DATASET_BASE = "/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/VLA/doanh/video_world/video_gen_physics/datasets/single_arm/multiview"
OUTPUT_BASE = "/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/VLA/doanh/video_world/video_gen_physics/models/Ctrl-World/dataset_example"

VIEW_ORDER = [
    "observation.images.exterior_1_left",
    "observation.images.exterior_2_left",
    "observation.images.wrist_left",
]

DOWNSAMPLE_FACTOR = 3  # 15fps -> 5fps


def find_episodes(dataset_dir):
    """Find all episode parquet files and return sorted list of (chunk_id, episode_id)."""
    data_dir = Path(dataset_dir) / "data"
    episodes = []
    for chunk_dir in sorted(data_dir.glob("chunk-*")):
        for pq in sorted(chunk_dir.glob("episode_*.parquet")):
            ep_id = int(pq.stem.split("_")[1])
            chunk_id = int(chunk_dir.name.split("-")[1])
            episodes.append((chunk_id, ep_id))
    return episodes


def get_video_path(dataset_dir, chunk_id, view_name, episode_id):
    return Path(dataset_dir) / "videos" / f"chunk-{chunk_id:03d}" / view_name / f"episode_{episode_id:06d}.mp4"


def downsample_and_resize_video(input_path, output_path, downsample=3, target_h=192, target_w=320):
    """Use ffmpeg fps filter to downsample 15fps->5fps and resize video."""
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    target_fps = 15 // downsample  # 5fps
    cmd = [
        "ffmpeg", "-y", "-i", str(input_path),
        "-vf", f"fps={target_fps},scale={target_w}:{target_h}",
        "-c:v", "libx264", "-preset", "fast", "-crf", "18",
        "-an", str(output_path),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"  WARNING: ffmpeg failed for {input_path}: {result.stderr[:200]}")
        return False
    return True


def get_frame_count(video_path):
    """Get frame count of a video using ffprobe."""
    cmd = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", str(video_path)]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        return 0
    info = json.loads(result.stdout)
    return int(info["streams"][0].get("nb_frames", 0))


def process_episode(dataset_dir, subset_name, chunk_id, episode_id, output_dir, input_save_dir=None):
    """Process one episode: read parquet, downsample videos, write annotation JSON."""
    parquet_path = Path(dataset_dir) / "data" / f"chunk-{chunk_id:03d}" / f"episode_{episode_id:06d}.parquet"
    if not parquet_path.exists():
        return None

    df = pd.read_parquet(parquet_path)
    n_frames_raw = len(df)

    # Downsample parquet data (every 3rd frame)
    df_ds = df.iloc[::DOWNSAMPLE_FACTOR].reset_index(drop=True)
    n_frames = len(df_ds)

    if n_frames < 10:
        return None

    # Extract states: cartesian_position (6D) + gripper (1D) = 7D
    cart_pos = np.array(df_ds["observation.state.cartesian_position"].tolist())  # (N, 6)
    gripper = np.array(df_ds["observation.state.gripper_position"].tolist())  # (N,) scalar
    if gripper.ndim == 1:
        gripper = gripper.reshape(-1, 1)
    states = np.concatenate([cart_pos, gripper], axis=1).tolist()  # (N, 7)

    # Extract joints: joint_position (7D) + gripper (1D) = 8D
    joint_pos = np.array(df_ds["observation.state.joint_position"].tolist())  # (N, 7)
    joints = np.concatenate([joint_pos, gripper], axis=1).tolist()  # (N, 8)

    # Language instruction
    instruction = df_ds["language_instruction"].iloc[0]
    if not instruction or instruction == "":
        instruction = "robot manipulation task"

    # Episode ID string for Ctrl-World
    ep_id_str = f"{episode_id:06d}"

    # Process videos
    video_entries = []
    for view_idx, view_name in enumerate(VIEW_ORDER):
        src_video = get_video_path(dataset_dir, chunk_id, view_name, episode_id)
        if not src_video.exists():
            print(f"  Missing video: {src_video}")
            return None

        dst_video = Path(output_dir) / subset_name / "videos" / "val" / ep_id_str / f"{view_idx}.mp4"
        if not dst_video.exists():
            success = downsample_and_resize_video(src_video, dst_video)
            if not success:
                return None

        rel_path = f"videos/val/{ep_id_str}/{view_idx}.mp4"
        video_entries.append({"video_path": rel_path})

    # Verify video frame count matches downsampled parquet
    first_video = Path(output_dir) / subset_name / "videos" / "val" / ep_id_str / "0.mp4"
    video_frames = get_frame_count(first_video)

    # Use the minimum of parquet-derived and video-derived frame count
    actual_length = min(n_frames, video_frames) if video_frames > 0 else n_frames
    states = states[:actual_length]
    joints = joints[:actual_length]

    # Build annotation
    annotation = {
        "texts": [instruction],
        "episode_id": episode_id,
        "success": True,
        "video_length": actual_length,
        "state_length": actual_length,
        "raw_length": n_frames_raw,
        "videos": video_entries,
        "states": states,
        "joints": joints,
    }

    # Save annotation
    anno_dir = Path(output_dir) / subset_name / "annotation" / "val"
    anno_dir.mkdir(parents=True, exist_ok=True)
    anno_path = anno_dir / f"{ep_id_str}.json"
    with open(anno_path, "w") as f:
        json.dump(annotation, f)

    # Save input for benchmark (GT videos + metadata)
    if input_save_dir is not None:
        import shutil
        ep_input_dir = Path(input_save_dir) / f"episode_{ep_id_str}"
        ep_input_dir.mkdir(parents=True, exist_ok=True)

        # Copy downsampled videos (all 3 views)
        for view_idx in range(len(VIEW_ORDER)):
            src = Path(output_dir) / subset_name / "videos" / "val" / ep_id_str / f"{view_idx}.mp4"
            dst = ep_input_dir / f"view_{view_idx}.mp4"
            if src.exists() and not dst.exists():
                shutil.copy2(str(src), str(dst))

        # Save metadata
        input_meta = {
            "episode_id": episode_id,
            "instruction": instruction,
            "num_frames": actual_length,
            "raw_frames": n_frames_raw,
            "fps": 5,
            "resolution": "320x192",
            "states": states,
            "joints": joints,
        }
        meta_path = ep_input_dir / "metadata.json"
        if not meta_path.exists():
            with open(meta_path, "w") as f:
                json.dump(input_meta, f, indent=2)

    return {"id": ep_id_str, "length": actual_length, "instruction": instruction}


INPUT_SAVE_BASE = "/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/VLA/doanh/video_world/video_gen_physics/sampling_dataset/dense/single_arm/input/multiview/ctrlworld"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--subset", choices=["makovian", "non_makovian", "both"], default="both")
    parser.add_argument("--output-dir", type=str, default=OUTPUT_BASE)
    parser.add_argument("--dataset-base", type=str, default=DATASET_BASE)
    parser.add_argument("--input-save-dir", type=str, default=INPUT_SAVE_BASE)
    args = parser.parse_args()

    subsets = ["makovian", "non_makovian"] if args.subset == "both" else [args.subset]

    for subset in subsets:
        dataset_dir = os.path.join(args.dataset_base, subset)
        subset_output_name = f"single_arm_multiview_{subset}"
        subset_input_save_dir = os.path.join(args.input_save_dir, subset)

        print(f"\n{'='*60}")
        print(f"Processing: {subset} -> {subset_output_name}")
        print(f"Input save: {subset_input_save_dir}")
        print(f"{'='*60}")

        episodes = find_episodes(dataset_dir)
        print(f"Found {len(episodes)} episodes")

        results = []
        for i, (chunk_id, ep_id) in enumerate(episodes):
            print(f"  [{i+1}/{len(episodes)}] chunk={chunk_id:03d} episode={ep_id:06d}", end=" ")
            result = process_episode(
                dataset_dir, subset_output_name, chunk_id, ep_id,
                args.output_dir, input_save_dir=subset_input_save_dir,
            )
            if result:
                results.append(result)
                print(f"-> {result['length']} frames")
            else:
                print("-> SKIPPED")

        # Save summary
        summary_path = os.path.join(args.output_dir, subset_output_name, "preparation_summary.json")
        os.makedirs(os.path.dirname(summary_path), exist_ok=True)
        with open(summary_path, "w") as f:
            json.dump({
                "total_episodes": len(results),
                "episodes": results,
            }, f, indent=2)

        print(f"\nDone: {len(results)} episodes prepared for {subset}")


if __name__ == "__main__":
    main()