| import os |
| import json |
| from collections import defaultdict |
|
|
| import cv2 |
| import h5py |
| import numpy as np |
| from dataset_upload.helpers import generate_unique_id |
|
|
|
|
| DEBUG = False |
|
|
|
|
| class Ph2dFrameloader: |
| """Pickle-able loader to read frames from an HDF5 file on demand. |
| |
| Reads frames from either 'observation.image.right' or 'observation.image.left'. |
| Each stored frame is expected to be an encoded image buffer which is decoded |
| with cv2.imdecode and converted to RGB. |
| """ |
|
|
| def __init__(self, hdf5_path: str, camera: str = "right") -> None: |
| self.hdf5_path = hdf5_path |
| self.camera = camera.lower() |
|
|
| def __call__(self) -> np.ndarray: |
| key = "observation.image.right" if self.camera == "right" else "observation.image.left" |
|
|
| with h5py.File(self.hdf5_path, "r") as f: |
| if key not in f: |
| raise KeyError(f"Key '{key}' not found in {self.hdf5_path}") |
|
|
| data = f[key] |
| encoded_frames = data[()] |
| frames = [] |
| for encoded_frame in encoded_frames: |
| frame = cv2.imdecode(encoded_frame, cv2.IMREAD_COLOR) |
| frames.append(frame) |
|
|
| frames_np = np.asarray(frames, dtype=np.uint8) |
| if frames_np.ndim != 4 or frames_np.shape[-1] != 3: |
| raise ValueError(f"Unexpected frames shape for {self.hdf5_path}: {frames_np.shape} (expected (T,H,W,3))") |
| return frames_np |
|
|
|
|
| def create_new_trajectory(hdf5_path: str, is_robot: bool, caption: str, camera: str) -> dict: |
| trajectory_info = {} |
| trajectory_info["id"] = generate_unique_id() |
| trajectory_info["task"] = caption |
| trajectory_info["frames"] = Ph2dFrameloader(hdf5_path, camera=camera) |
| trajectory_info["is_robot"] = is_robot |
| trajectory_info["quality_label"] = "successful" |
| trajectory_info["partial_success"] = 1 |
| trajectory_info["data_source"] = "ph2d" |
| return trajectory_info |
|
|
|
|
| def load_ph2d_dataset(dataset_path: str) -> dict[str, list[dict]]: |
| """Load Ph2d dataset organized as folders with HDF5 files and optional metadata.json. |
| |
| Expected layout: |
| <dataset_path>/ |
| ph2d_metadata.json # Optional (dataset-specific). TODO: parse to captions/tasks |
| <sequence_1>/ |
| *.hdf5 |
| <sequence_2>/ |
| ... |
| |
| Args: |
| dataset_path: Root directory containing sequence folders of HDF5 files. |
| |
| Returns: |
| Mapping: task/caption -> list of trajectory dicts. |
| """ |
|
|
| task_data: dict[str, list[dict]] = defaultdict(list) |
|
|
| |
| all_task_attributes = json.load(open(os.path.join(dataset_path, "ph2d_metadata.json")))["per_task_attributes"] |
| for task_name, task_attributes in all_task_attributes.items(): |
| |
| embodiment_type = task_attributes["embodiment_type"] |
| if "human" in embodiment_type: |
| is_robot = False |
| else: |
| is_robot = True |
|
|
| for h5_file in os.listdir(os.path.join(dataset_path, task_name)): |
| if not h5_file.endswith(".hdf5"): |
| continue |
|
|
| h5_path = os.path.join(dataset_path, task_name, h5_file) |
| loaded_data = h5py.File(h5_path, "r") |
| task_description = loaded_data.attrs["description"] |
|
|
| if "observation.image.right" in loaded_data: |
| task_data[task_name].append(create_new_trajectory(h5_path, is_robot, task_description, camera="right")) |
| if "observation.image.left" in loaded_data: |
| task_data[task_name].append(create_new_trajectory(h5_path, is_robot, task_description, camera="left")) |
|
|
| if DEBUG: |
| break |
| print(f"Loaded {sum(len(task_list) for task_list in task_data.values())} trajectories from {len(task_data)} tasks") |
| return task_data |
|
|
