File size: 5,547 Bytes
319eb16 | 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 | #!/usr/bin/env python3
from __future__ import annotations
from collections import defaultdict
from pathlib import Path
from typing import Iterable
import numpy as np
import cv2
from dataset_upload.helpers import generate_unique_id
TASK_DESCRIPTION_MAP = {
"Add_pepper_to_the_green_bowl": "Add pepper to the green bowl",
"Collect_the_fork_to_the_yellow_box": "Collect the fork to the yellow box",
"Press_the_button": "Press the button",
"Put_the_bread_in_the_oven": "Put the bread in the oven",
"Put_the_fruit_in_the_yellow_plate": "Put the apple on the yellow plate",
"Put_the_marker_into_the_pen_cup": "Put the marker into the pen cup",
"Put_the_red_bowl_on_the_blue_plate": "Put the red bowl on the blue plate",
"Put_the_red_cup_on_the_purple_coaster": "Put the red cup on the purple coaster",
"Put_the_rubber_to_the_blue_pencil_box": "Put the eraser in the blue pencil box",
"Stack_the_green_block_on_the_red_block": "Stack the green block on the red block",
}
QUALITY_LABEL_MAP = {
"succ": "successful",
"success": "successful",
"successful": "successful",
"subopt": "suboptimal",
"suboptimal": "suboptimal",
"fail": "failure",
"failure": "failure",
}
class UTDSO101FrameLoader:
"""Lazy loader that extracts RGB frames from MP4 video files."""
def __init__(self, video_path: str) -> None:
if not Path(video_path).exists():
raise FileNotFoundError(f"Video file not found: {video_path}")
self.video_path = video_path
def __call__(self) -> np.ndarray:
"""Load all frames from the MP4 video file."""
cap = cv2.VideoCapture(self.video_path)
if not cap.isOpened():
raise ValueError(f"Could not open video file: {self.video_path}")
frames: list[np.ndarray] = []
while True:
ret, frame_bgr = cap.read()
if not ret:
break
# Convert BGR to RGB
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
frames.append(frame_rgb)
cap.release()
if not frames:
raise ValueError(f"No frames found in video: {self.video_path}")
return np.stack(frames, axis=0).astype(np.uint8)
def _default_task_description(task_key: str) -> str:
"""Convert task key to a readable description."""
if task_key in TASK_DESCRIPTION_MAP:
return TASK_DESCRIPTION_MAP[task_key]
return task_key.replace("_", " ").capitalize()
def _parse_video_metadata(video_filename: str) -> tuple[str, str, str]:
"""Parse video filename to extract task, optimality, and demo_idx.
Expected format: {task}_{optimality}_{demo_idx}.mp4
Example: pick_blue_cup_success_1.mp4
"""
# Remove .mp4 extension
name_without_ext = video_filename.replace(".mp4", "")
parts = name_without_ext.split("_")
if len(parts) < 3:
raise ValueError(f"Unexpected video filename format: {video_filename}")
demo_idx = parts[-1]
optimality_key = parts[-2].lower()
task_key = "_".join(parts[:-2])
return task_key, optimality_key, demo_idx
def load_utd_so101_dataset(
dataset_path: str, max_trajectories: int | None = None, is_robot: bool = True, data_source: str = "utd_so101"
) -> dict[str, list[dict]]:
"""Load UTD SO101 trajectories grouped by language task.
Args:
dataset_path: Path to the dataset directory
max_trajectories: Maximum number of trajectories to load
is_robot: Whether trajectories are robot (True) or human (False)
data_source: Data source identifier for the dataset
"""
root = Path(dataset_path).expanduser()
if not root.exists():
raise FileNotFoundError(f"Dataset path not found: {dataset_path}")
if (root / "koch_arm_ut_dallas").exists():
root = root / "koch_arm_ut_dallas"
# Find all MP4 files
video_files = sorted([p for p in root.glob("*.mp4")])
if not video_files:
raise ValueError(f"No MP4 files found in {root}")
limit = None if max_trajectories is None or max_trajectories < 0 else int(max_trajectories)
task_data: dict[str, list[dict]] = defaultdict(list)
total = 0
for video_path in video_files:
if limit is not None and total >= limit:
break
try:
task_key, optimality_key, demo_idx = _parse_video_metadata(video_path.name)
except ValueError as e:
print(f"⚠️ Skipping {video_path.name}: {e}")
continue
# For human videos, always set quality_label to "successful"
if not is_robot:
quality_label = "successful"
elif optimality_key not in QUALITY_LABEL_MAP:
print(f"⚠️ Skipping {video_path.name}: Unknown optimality label '{optimality_key}'")
continue
else:
quality_label = QUALITY_LABEL_MAP[optimality_key]
frame_loader = UTDSO101FrameLoader(str(video_path))
task_description = _default_task_description(task_key)
trajectory = {
"id": generate_unique_id(),
"task": task_description,
"frames": frame_loader,
"is_robot": is_robot,
"quality_label": quality_label,
"data_source": data_source,
}
task_data[task_description].append(trajectory)
total += 1
dataset_type = "robot" if is_robot else "human"
print(f"Loaded {total} {dataset_type} trajectories from {len(task_data)} tasks in UTD SO101 dataset")
return task_data
|