add support for robomind2lerobot (#26)

Co-authored-by: HaomingSong <haomingsong24@gmail.com>
2026-07-25 13:05:57 +00:00 · 2025-05-12 20:16:13 +08:00
parent 37f856f50b
commit 7af65ba23a
16 changed files with 1445 additions and 1 deletions
@@ -0,0 +1,74 @@
+from pathlib import Path
+
+import cv2
+import h5py
+import numpy as np
+
+
+def decode_images(camera_key, input_images):
+    if "depth" not in camera_key:
+        rgb_images = []
+        camera_rgb_images = input_images
+        for camera_rgb_image in camera_rgb_images:
+            camera_rgb_image = np.array(camera_rgb_image)
+            rgb = cv2.imdecode(camera_rgb_image, cv2.IMREAD_COLOR)
+            if rgb is None:
+                rgb = np.frombuffer(camera_rgb_image, dtype=np.uint8)
+                if rgb.size == 2764800:
+                    rgb = rgb.reshape(720, 1280, 3)
+                elif rgb.size == 921600:
+                    rgb = rgb.reshape(480, 640, 3)
+            rgb_images.append(rgb)
+        rgb_images = np.asarray(rgb_images)
+        return rgb_images
+    else:
+        depth_images = []
+        camera_depth_images = input_images
+        for camera_depth_image in camera_depth_images:
+            if isinstance(camera_depth_image, np.ndarray):
+                depth_array = camera_depth_image
+            else:
+                depth_array = np.frombuffer(camera_depth_image, dtype=np.uint8)
+            depth = cv2.imdecode(depth_array, cv2.IMREAD_UNCHANGED)
+            if depth is None:
+                if depth_array.size == 921600:
+                    depth = depth_array.reshape(720, 1280)
+                elif depth_array.size == 307200:
+                    depth = depth_array.reshape(480, 640)
+            depth_images.append(depth)
+        depth_images = np.asarray(depth_images)[..., None]
+        return depth_images
+
+
+def load_local_dataset(episode_path: Path, config: dict, save_depth: bool):
+    try:
+        images = {}
+        states = {}
+        actions = {}
+        with h5py.File(episode_path, "r") as file:
+            for key in config["images"]:
+                if save_depth and "depth" in key:
+                    image_key = f"observations/depth_images/{key[:-6]}"
+                elif "depth" not in key:
+                    image_key = f"observations/rgb_images/{key}"
+                else:
+                    continue
+                images[f"observation.images.{key}"] = decode_images(image_key, file[image_key])
+            for key in config["states"]:
+                states[f"observation.states.{key}"] = np.array(file[f"puppet/{key}"], dtype=np.float32)
+            for key in config["actions"]:
+                actions[f"actions.{key}"] = np.array(file[f"master/{key}"], dtype=np.float32)
+
+        num_frames = len(next(iter(states.values())))
+        frames = [
+            {
+                **{key: value[i] for key, value in images.items() if save_depth or "depth" not in key},
+                **{key: value[i] for key, value in states.items()},
+                **{key: value[i] for key, value in actions.items()},
+            }
+            for i in range(num_frames)
+        ]
+        return True, frames, ""
+
+    except (FileNotFoundError, OSError, KeyError) as e:
+        return False, [], e