File size: 5,864 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 160 161 162 | #!/usr/bin/env python3
"""
Metaworld dataset loader for the generic dataset converter for Robometer model training.
This module contains logic for loading metaworld data organized by task and quality.
uv run python dataset_upload/generate_hf_dataset.py \
--config_path=dataset_upload/configs/data_gen_configs/metaworld.yaml
"""
import collections
from pathlib import Path
import h5py
import numpy as np
from PIL import Image
from torchvision import transforms
from tqdm import tqdm
from dataset_upload.video_helpers import load_video_frames
from dataset_upload.dataset_loaders.mw_task_annotations import TRAIN_GT_ANN, EVAL_GT_ANN
def apply_center_crop_to_frames(frames: np.ndarray) -> np.ndarray:
"""Apply center crop (224, 224) to video frames using torchvision transforms.
Args:
frames: numpy array of shape (T, H, W, 3) in RGB order
Returns:
numpy array of shape (T, 224, 224, 3) with center cropped frames
"""
# Define the center crop transform
center_crop = transforms.CenterCrop(224)
cropped_frames = []
for frame in frames:
# Convert numpy array to PIL Image
pil_frame = Image.fromarray(frame.astype(np.uint8))
# Apply center crop
cropped_pil = center_crop(pil_frame)
# Convert back to numpy array
cropped_frame = np.array(cropped_pil)
cropped_frames.append(cropped_frame)
return np.array(cropped_frames)
def map_quality_label(original_label: str) -> str:
"""Map original quality labels to standardized Robometer labels."""
label_mapping = {"GT": "successful", "success": "successful", "all_fail": "failure", "close_succ": "suboptimal"}
return label_mapping.get(original_label, original_label)
def load_metaworld_dataset(base_path: str, dataset_name: str) -> dict[str, list[dict]]:
"""Load metaworld dataset and organize by task.
Args:
base_path: Path to the metaworld dataset directory
Returns:
Dictionary mapping task names to lists of trajectory dictionaries
"""
print(f"Loading metaworld dataset from: {base_path}")
task_data = collections.defaultdict(list)
base_path = Path(base_path)
if not base_path.exists():
raise FileNotFoundError(f"Metaworld dataset path not found: {base_path}")
if "train" in dataset_name:
tasks = TRAIN_GT_ANN.keys()
anns = TRAIN_GT_ANN
elif "eval" in dataset_name:
tasks = EVAL_GT_ANN.keys()
anns = EVAL_GT_ANN
print("number of tasks: ", len(tasks))
scucessful_trajs_file = "downloaded_data/metaworld_video.h5"
with h5py.File(scucessful_trajs_file, "r") as f:
print("Available tasks: ", f.keys())
print("number of tasks: ", len(f.keys()))
task_names = list(f.keys())
for task_name in tqdm(task_names, desc="Loading successful trajectories"):
if task_name not in tasks:
continue
for traj_name in ["0", "1", "10", "11", "12"]: # not sure why we use these, but we use these 5
traj = f[task_name][traj_name]
frames = np.array(traj)
cropped_frames = apply_center_crop_to_frames(frames)
trajectory = {
"frames": cropped_frames,
"task": anns.get(task_name),
"quality_label": "successful",
"is_robot": True,
"partial_success": 0,
}
task_data[task_name].append(trajectory)
if "eval" in dataset_name:
task_dirs = [d for d in base_path.iterdir() if d.is_dir() and not d.name.startswith(".")]
for task_dir in tqdm(task_dirs, desc="Processing tasks"):
task_name = task_dir.name
if task_name in [".DS_Store"]:
continue
# Find all quality label directories within this task
quality_dirs = [d for d in task_dir.iterdir() if d.is_dir() and not d.name.startswith(".")]
for quality_dir in quality_dirs:
original_quality_label = quality_dir.name
# Map quality label to standardized format
quality_label = map_quality_label(original_quality_label)
# Find all video files in this quality directory
video_files = list(quality_dir.glob("*.mp4")) + list(quality_dir.glob("*.gif"))
# if len(video_files) != 2:
# import ipdb; ipdb.set_trace()
for video_file in video_files:
# # Extract index from filename (e.g., "1.mp4" -> 1)
# try:
# int(video_file.stem)
# except ValueError:
# print(f"Warning: Could not parse index from filename: {video_file.name}")
# Load frames and apply center crop
original_frames = load_video_frames(video_file)
cropped_frames = apply_center_crop_to_frames(original_frames)
nl_name = anns.get(task_name)
# Create trajectory entry
trajectory = {
"frames": cropped_frames,
"task": nl_name, # Natural language task
"quality_label": quality_label, # Mapped quality label
"is_robot": True,
"partial_success": 0,
}
task_data[task_name].append(trajectory)
for task_name, trajectories in task_data.items():
print(f"Task {task_name}: {len(trajectories)} trajectories")
total_trajectories = sum(len(trajectories) for trajectories in task_data.values())
print(f"\nLoaded {total_trajectories} trajectories from {len(task_data)} tasks")
return task_data
|