File size: 7,061 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 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | import json
import os
from pathlib import Path
from typing import Any
import cv2
import numpy as np
from dataset_upload.helpers import generate_unique_id
CAMERA_DIR_CANDIDATES = [
"front_rgb",
"left_shoulder_rgb",
"right_shoulder_rgb",
# "right_shoudler_rgb", # sometimes misspelled in datasets
# "wrist_rgb",
]
class RacerFrameListLoader:
"""Pickle-able loader that reads a list of image paths on demand (RGB, uint8)."""
def __init__(self, image_paths: list[str]) -> None:
if not image_paths:
raise ValueError("image_paths must be non-empty")
self.image_paths = image_paths
def __call__(self) -> np.ndarray:
frames: list[np.ndarray] = []
for p in self.image_paths:
img_bgr = cv2.imread(p, cv2.IMREAD_COLOR)
if img_bgr is None:
continue
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
frames.append(img_rgb)
if not frames:
return np.empty((0, 0, 0, 3), dtype=np.uint8)
frames_np = np.asarray(frames, dtype=np.uint8)
return frames_np
def _sorted_pngs(dir_path: Path) -> list[str]:
paths = [p for p in dir_path.glob("*.png")]
paths.sort(key=lambda x: int(x.stem.split("_")[0]))
return [str(p) for p in paths]
def _make_traj(image_paths: list[str], task_text: str, is_success: bool) -> dict[str, Any]:
traj: dict[str, Any] = {}
traj["id"] = generate_unique_id()
traj["task"] = task_text
traj["frames"] = RacerFrameListLoader(image_paths)
traj["is_robot"] = True
traj["quality_label"] = "successful" if is_success else "failure"
traj["data_source"] = "racer"
traj["preference_group_id"] = None
traj["preference_rank"] = None
return traj
def _collect_camera_views(sample_dir: Path) -> dict[str, list[str]]:
views: dict[str, list[str]] = {}
for cam in CAMERA_DIR_CANDIDATES:
d = sample_dir / cam
if d.exists() and d.is_dir():
imgs = _sorted_pngs(d)
if imgs:
views[cam] = imgs
return views
def load_racer_dataset(dataset_path: str, dataset_name: str) -> dict[str, list[dict]]:
"""Load RACER-augmented_rlbench dataset.
Args:
dataset_path: Path to dataset root containing 'train' and/or 'val' folders.
dataset_name: Use to pick split: 'racer_train' -> train, 'racer_val' -> val.
Behavior:
- Uses task_goal from language_description.json as language instruction.
- Creates success trajectories (full expert episode) per camera view.
- For each expert subgoal frame that contains heuristic failures in 'augmentation',
creates failure trajectories up to that expert frame index (inclusive), per camera view.
Returns:
Mapping: task_goal -> list of trajectory dicts.
"""
root = Path(os.path.expanduser(dataset_path))
if not root.exists():
raise FileNotFoundError(f"RACER dataset path not found: {root}")
split = "val" if ("val" in dataset_name.lower()) else "train"
# Some distributions include an extra 'samples' segment
split_dir = root / split
alt_split_dir = split_dir / "samples"
if alt_split_dir.exists():
split_dir = alt_split_dir
if not split_dir.exists():
raise FileNotFoundError(f"Split directory not found: {split_dir}")
# Tasks are subdirectories under split_dir
task_dirs = [p for p in split_dir.iterdir() if p.is_dir()]
task_data: dict[str, list[dict]] = {}
for task_dir in task_dirs:
# Episodes are numeric directories under each task
episode_dirs = [p for p in task_dir.iterdir() if p.is_dir()]
for ep_dir in episode_dirs:
json_path = ep_dir / "language_description.json"
if not json_path.exists():
continue
try:
with open(json_path, "r") as f:
desc = json.load(f)
except Exception:
continue
task_goal: str = desc.get("task_goal", "").strip() or task_dir.name
subgoal_dict: dict[str, Any] = desc.get("subgoal", {}) or {}
# Gather camera views for this episode once
views = _collect_camera_views(ep_dir)
if not views:
continue
# Success: use full length per view
for cam, img_list in views.items():
if not img_list:
continue
expert_img_list = [p for p in img_list if "expert" in p]
traj = _make_traj(expert_img_list, task_goal, is_success=True)
task_data.setdefault(task_goal, []).append(traj)
# Failures: for each expert key that contains heuristic augmentations
for key, content in subgoal_dict.items():
# Expect keys like '0_expert', '48_expert', ...
if not isinstance(key, str) or "expert" not in key:
continue
try:
expert_frame_idx = int(key.split("_")[0])
except Exception:
continue
aug = content.get("augmentation", {}) if isinstance(content, dict) else {}
if not isinstance(aug, dict) or not aug:
continue
# Enumerate augmentations; select those labeled as heuristic failures
has_failure = False
for aug_image_name, aug_content in aug.items():
if not isinstance(aug_content, dict):
continue
label = str(aug_content.get("label", "")).lower()
if "failure" in label: # e.g., 'recoverable_failure'
has_failure = True
break
if not has_failure:
continue
# Build failure trajectories by truncating expert frames up to expert_frame_idx
for cam, img_list in views.items():
if not img_list:
continue
# Find frames with numeric names and truncate accordingly
def _frame_num(p: str) -> int:
try:
return int(Path(p).stem.split("_")[0])
except Exception:
return 1_000_000_000
# Keep frames with index < expert_frame_idx
subset = [p for p in img_list if _frame_num(p) < expert_frame_idx and "expert" in p]
# add the augmented failure frame
for img_name in img_list:
if aug_image_name in img_name:
subset.append(img_name)
break
if not subset:
continue
traj = _make_traj(subset, task_goal, is_success=False)
task_data.setdefault(task_goal, []).append(traj)
return task_data
|