mirror of
https://github.com/Tavish9/any4lerobot.git
synced 2026-07-25 21:15:59 +00:00
add support for robomind2lerobot (#26)
Co-authored-by: HaomingSong <haomingsong24@gmail.com>
This commit is contained in:
Qizhi Chen
@@ -0,0 +1,21 @@
|
||||
from .agilex_3rgb import AgileX_3RGB_Config
|
||||
from .franka_1rgb import Franka_1RGB_Config
|
||||
from .franka_3rgb import Franka_3RGB_Config
|
||||
from .franka_fr3_dual_arm import Franka_Fr3_Dual_Arm_Config
|
||||
from .tienkung_gello_1rgb import Tien_Kung_Gello_1RGB_Config
|
||||
from .tienkung_prod1_gello_1rgb import Tien_Kung_Prod1_Gello_1RGB_Config
|
||||
from .tienkung_xsens_1rgb import Tien_Kung_Xsens_1RGB_Config
|
||||
from .ur_1rgb import UR_1RGB_Config
|
||||
|
||||
ROBOMIND_CONFIG = dict(
|
||||
agilex_3rgb=AgileX_3RGB_Config,
|
||||
franka_1rgb=Franka_1RGB_Config,
|
||||
franka_3rgb=Franka_3RGB_Config,
|
||||
franka_fr3_dual=Franka_Fr3_Dual_Arm_Config,
|
||||
sim_franka_3rgb="",
|
||||
sim_tienkung_1rgb="",
|
||||
tienkung_gello_1rgb=Tien_Kung_Gello_1RGB_Config,
|
||||
tienkung_prod1_gello_1rgb=Tien_Kung_Prod1_Gello_1RGB_Config,
|
||||
tienkung_xsens_1rgb=Tien_Kung_Xsens_1RGB_Config,
|
||||
ur_1rgb=UR_1RGB_Config,
|
||||
)
|
||||
@@ -0,0 +1,118 @@
|
||||
AgileX_3RGB_Config = {
|
||||
"images": {
|
||||
"camera_front": {
|
||||
"dtype": "video",
|
||||
"shape": (480, 640, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_left_wrist": {
|
||||
"dtype": "video",
|
||||
"shape": (480, 640, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_right_wrist": {
|
||||
"dtype": "video",
|
||||
"shape": (480, 640, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_front_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
"camera_left_wrist_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
"camera_right_wrist_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
},
|
||||
"states": {
|
||||
"end_effector_left": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["x", "y", "z", "rx", "ry", "rz", "rw"]},
|
||||
},
|
||||
"end_effector_right": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["x", "y", "z", "rx", "ry", "rz", "rw"]},
|
||||
},
|
||||
"joint_effort_left": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
"joint_effort_right": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
"joint_position_left": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
"joint_position_right": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
"joint_velocity_left": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
"joint_velocity_right": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
},
|
||||
"actions": {
|
||||
"end_effector_left": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["x", "y", "z", "rx", "ry", "rz", "rw"]},
|
||||
},
|
||||
"end_effector_right": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["x", "y", "z", "rx", "ry", "rz", "rw"]},
|
||||
},
|
||||
"joint_effort_left": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
"joint_effort_right": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
"joint_position_left": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
"joint_position_right": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
"joint_velocity_left": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
"joint_velocity_right": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
Franka_1RGB_Config = {
|
||||
"images": {
|
||||
"camera_top": {
|
||||
"dtype": "video",
|
||||
"shape": (720, 1280, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_top_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
},
|
||||
"states": {
|
||||
"end_effector": {
|
||||
"dtype": "float32",
|
||||
"shape": (6,),
|
||||
"names": {"motors": ["x", "y", "z", "r", "p", "y"]},
|
||||
},
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (8,),
|
||||
"names": {
|
||||
"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6", "gripper"]
|
||||
},
|
||||
},
|
||||
},
|
||||
"actions": {
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (8,),
|
||||
"names": {
|
||||
"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6", "gripper"]
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
Franka_3RGB_Config = {
|
||||
"images": {
|
||||
"camera_top": {
|
||||
"dtype": "video",
|
||||
"shape": (720, 1280, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_left": {
|
||||
"dtype": "video",
|
||||
"shape": (480, 640, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_right": {
|
||||
"dtype": "video",
|
||||
"shape": (480, 640, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_top_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (720, 1280, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
"camera_left_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
"camera_right_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
},
|
||||
"states": {
|
||||
"end_effector": {
|
||||
"dtype": "float32",
|
||||
"shape": (6,),
|
||||
"names": {"motors": ["x", "y", "z", "r", "p", "y"]},
|
||||
},
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (8,),
|
||||
"names": {
|
||||
"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6", "gripper"]
|
||||
},
|
||||
},
|
||||
},
|
||||
"actions": {
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (8,),
|
||||
"names": {
|
||||
"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6", "gripper"]
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
Franka_Fr3_Dual_Arm_Config = {
|
||||
"images": {
|
||||
"camera_front": {
|
||||
"dtype": "video",
|
||||
"shape": (720, 1280, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_top": {
|
||||
"dtype": "video",
|
||||
"shape": (720, 1280, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_left": {
|
||||
"dtype": "video",
|
||||
"shape": (480, 640, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_right": {
|
||||
"dtype": "video",
|
||||
"shape": (480, 640, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_front_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (720, 1280, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
"camera_top_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (720, 1280, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
"camera_left_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
"camera_right_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
},
|
||||
"states": {
|
||||
"end_effector": {
|
||||
"dtype": "float32",
|
||||
"shape": (12,),
|
||||
"names": {"motors": ["left_xyzrpy", "right_xyzrpy"]},
|
||||
},
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (16,),
|
||||
"names": {
|
||||
"motors": [
|
||||
"left_joint_0",
|
||||
"left_joint_1",
|
||||
"left_joint_2",
|
||||
"left_joint_3",
|
||||
"left_joint_4",
|
||||
"left_joint_5",
|
||||
"left_joint_6",
|
||||
"left_gripper",
|
||||
"right_joint_0",
|
||||
"right_joint_1",
|
||||
"right_joint_2",
|
||||
"right_joint_3",
|
||||
"right_joint_4",
|
||||
"right_joint_5",
|
||||
"right_joint_6",
|
||||
"right_gripper",
|
||||
]
|
||||
},
|
||||
},
|
||||
},
|
||||
"actions": {
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (16,),
|
||||
"names": {
|
||||
"motors": [
|
||||
"left_joint_0",
|
||||
"left_joint_1",
|
||||
"left_joint_2",
|
||||
"left_joint_3",
|
||||
"left_joint_4",
|
||||
"left_joint_5",
|
||||
"left_joint_6",
|
||||
"left_gripper",
|
||||
"right_joint_0",
|
||||
"right_joint_1",
|
||||
"right_joint_2",
|
||||
"right_joint_3",
|
||||
"right_joint_4",
|
||||
"right_joint_5",
|
||||
"right_joint_6",
|
||||
"right_gripper",
|
||||
]
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
Tien_Kung_Gello_1RGB_Config = {
|
||||
"images": {
|
||||
"camera_top": {
|
||||
"dtype": "video",
|
||||
"shape": (480, 640, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_top_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
},
|
||||
"states": {
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (16,),
|
||||
"names": {
|
||||
"motors": [
|
||||
"left_arm_0",
|
||||
"left_arm_1",
|
||||
"left_arm_2",
|
||||
"left_arm_3",
|
||||
"left_arm_4",
|
||||
"left_arm_5",
|
||||
"left_arm_6",
|
||||
"left hand_closure",
|
||||
"right_arm_0",
|
||||
"right_arm_1",
|
||||
"right_arm_2",
|
||||
"right_arm_3",
|
||||
"right_arm_4",
|
||||
"right_arm_5",
|
||||
"right_arm_6",
|
||||
"right hand closure",
|
||||
]
|
||||
},
|
||||
},
|
||||
},
|
||||
"actions": {
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (16,),
|
||||
"names": {
|
||||
"motors": [
|
||||
"left_arm_0",
|
||||
"left_arm_1",
|
||||
"left_arm_2",
|
||||
"left_arm_3",
|
||||
"left_arm_4",
|
||||
"left_arm_5",
|
||||
"left_arm_6",
|
||||
"left hand_closure",
|
||||
"right_arm_0",
|
||||
"right_arm_1",
|
||||
"right_arm_2",
|
||||
"right_arm_3",
|
||||
"right_arm_4",
|
||||
"right_arm_5",
|
||||
"right_arm_6",
|
||||
"right hand closure",
|
||||
]
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
Tien_Kung_Prod1_Gello_1RGB_Config = {
|
||||
"images": {
|
||||
"camera_top": {
|
||||
"dtype": "video",
|
||||
"shape": (720, 1280, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_top_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (720, 1280, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
},
|
||||
"states": {
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (16,),
|
||||
"names": {
|
||||
"motors": [
|
||||
"left_arm_0",
|
||||
"left_arm_1",
|
||||
"left_arm_2",
|
||||
"left_arm_3",
|
||||
"left_arm_4",
|
||||
"left_arm_5",
|
||||
"left_arm_6",
|
||||
"left hand_closure",
|
||||
"right_arm_0",
|
||||
"right_arm_1",
|
||||
"right_arm_2",
|
||||
"right_arm_3",
|
||||
"right_arm_4",
|
||||
"right_arm_5",
|
||||
"right_arm_6",
|
||||
"right hand closure",
|
||||
]
|
||||
},
|
||||
},
|
||||
},
|
||||
"actions": {
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (16,),
|
||||
"names": {
|
||||
"motors": [
|
||||
"left_arm_0",
|
||||
"left_arm_1",
|
||||
"left_arm_2",
|
||||
"left_arm_3",
|
||||
"left_arm_4",
|
||||
"left_arm_5",
|
||||
"left_arm_6",
|
||||
"left hand_closure",
|
||||
"right_arm_0",
|
||||
"right_arm_1",
|
||||
"right_arm_2",
|
||||
"right_arm_3",
|
||||
"right_arm_4",
|
||||
"right_arm_5",
|
||||
"right_arm_6",
|
||||
"right hand closure",
|
||||
]
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
Tien_Kung_Xsens_1RGB_Config = {
|
||||
"images": {
|
||||
"camera_top": {
|
||||
"dtype": "video",
|
||||
"shape": (480, 640, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_top_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
},
|
||||
"states": {
|
||||
"end_effector": {
|
||||
"dtype": "float32",
|
||||
"shape": (12,),
|
||||
"names": {
|
||||
"motors": [
|
||||
"left_little_finger",
|
||||
"left_ring_finger",
|
||||
"left_middle_finger",
|
||||
"left_index_finger",
|
||||
"left_thumb0_for_bending",
|
||||
"left_thumb1_for_rotation",
|
||||
"right_little_finger",
|
||||
"right_ring_finger",
|
||||
"right_middle_finger",
|
||||
"right_index_finger",
|
||||
"right_thumb0_for_bending",
|
||||
"right_thumb1_for_rotation",
|
||||
]
|
||||
},
|
||||
},
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (14,),
|
||||
"names": {
|
||||
"motors": [
|
||||
"left_arm_0",
|
||||
"left_arm_1",
|
||||
"left_arm_2",
|
||||
"left_arm_3",
|
||||
"left_arm_4",
|
||||
"left_arm_5",
|
||||
"left_arm_6",
|
||||
"right_arm_0",
|
||||
"right_arm_1",
|
||||
"right_arm_2",
|
||||
"right_arm_3",
|
||||
"right_arm_4",
|
||||
"right_arm_5",
|
||||
"right_arm_6",
|
||||
]
|
||||
},
|
||||
},
|
||||
},
|
||||
"actions": {
|
||||
"end_effector": {
|
||||
"dtype": "float32",
|
||||
"shape": (12,),
|
||||
"names": {
|
||||
"motors": [
|
||||
"left_little_finger",
|
||||
"left_ring_finger",
|
||||
"left_middle_finger",
|
||||
"left_index_finger",
|
||||
"left_thumb0_for_bending",
|
||||
"left_thumb1_for_rotation",
|
||||
"right_little_finger",
|
||||
"right_ring_finger",
|
||||
"right_middle_finger",
|
||||
"right_index_finger",
|
||||
"right_thumb0_for_bending",
|
||||
"right_thumb1_for_rotation",
|
||||
]
|
||||
},
|
||||
},
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (14,),
|
||||
"names": {
|
||||
"motors": [
|
||||
"left_arm_0",
|
||||
"left_arm_1",
|
||||
"left_arm_2",
|
||||
"left_arm_3",
|
||||
"left_arm_4",
|
||||
"left_arm_5",
|
||||
"left_arm_6",
|
||||
"right_arm_0",
|
||||
"right_arm_1",
|
||||
"right_arm_2",
|
||||
"right_arm_3",
|
||||
"right_arm_4",
|
||||
"right_arm_5",
|
||||
"right_arm_6",
|
||||
]
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
UR_1RGB_Config = {
|
||||
"images": {
|
||||
"camera_top": {
|
||||
"dtype": "video",
|
||||
"shape": (480, 640, 3),
|
||||
"names": ["height", "width", "rgb"],
|
||||
},
|
||||
"camera_top_depth": {
|
||||
"dtype": "image",
|
||||
"shape": (480, 640, 1),
|
||||
"names": ["height", "width", "channel"],
|
||||
},
|
||||
},
|
||||
"states": {
|
||||
"end_effector": {
|
||||
"dtype": "float32",
|
||||
"shape": (6,),
|
||||
"names": {"motors": ["x", "y", "z", "r", "p", "y"]},
|
||||
},
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
},
|
||||
"actions": {
|
||||
"joint_position": {
|
||||
"dtype": "float32",
|
||||
"shape": (7,),
|
||||
"names": {"motors": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
import numpy as np
|
||||
import torchvision
|
||||
from lerobot.common.datasets.compute_stats import auto_downsample_height_width, get_feature_stats, sample_indices
|
||||
from lerobot.common.datasets.utils import load_image_as_numpy
|
||||
|
||||
torchvision.set_video_backend("pyav")
|
||||
|
||||
|
||||
def generate_features_from_config(AgiBotWorld_CONFIG):
|
||||
features = {}
|
||||
for key, value in AgiBotWorld_CONFIG["images"].items():
|
||||
features[f"observation.images.{key}"] = value
|
||||
for key, value in AgiBotWorld_CONFIG["states"].items():
|
||||
features[f"observation.states.{key}"] = value
|
||||
for key, value in AgiBotWorld_CONFIG["actions"].items():
|
||||
features[f"actions.{key}"] = value
|
||||
return features
|
||||
|
||||
|
||||
def sample_images(input):
|
||||
if type(input) is list:
|
||||
image_paths = input
|
||||
|
||||
sampled_indices = sample_indices(len(image_paths))
|
||||
images = None
|
||||
for i, idx in enumerate(sampled_indices):
|
||||
path = image_paths[idx]
|
||||
|
||||
img = load_image_as_numpy(path, dtype=np.uint8, channel_first=True)
|
||||
img = auto_downsample_height_width(img)
|
||||
|
||||
if images is None:
|
||||
images = np.empty((len(sampled_indices), *img.shape), dtype=np.uint8)
|
||||
|
||||
images[i] = img
|
||||
elif type(input) is np.ndarray:
|
||||
frames_array = input[:, None, :, :] # Shape: [T, 1, H, W]
|
||||
sampled_indices = sample_indices(len(frames_array))
|
||||
images = None
|
||||
for i, idx in enumerate(sampled_indices):
|
||||
img = frames_array[idx]
|
||||
img = auto_downsample_height_width(img)
|
||||
|
||||
if images is None:
|
||||
images = np.empty((len(sampled_indices), *img.shape), dtype=np.uint8)
|
||||
|
||||
images[i] = img
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def compute_episode_stats(episode_data: dict[str, list[str] | np.ndarray], features: dict) -> dict:
|
||||
ep_stats = {}
|
||||
for key, data in episode_data.items():
|
||||
if features[key]["dtype"] == "string":
|
||||
continue # HACK: we should receive np.arrays of strings
|
||||
elif features[key]["dtype"] in ["image", "video"]:
|
||||
ep_ft_array = sample_images(data)
|
||||
axes_to_reduce = (0, 2, 3) # keep channel dim
|
||||
keepdims = True
|
||||
else:
|
||||
ep_ft_array = data # data is already a np.ndarray
|
||||
axes_to_reduce = 0 # compute stats over the first axis
|
||||
keepdims = data.ndim == 1 # keep as np.array
|
||||
|
||||
ep_stats[key] = get_feature_stats(ep_ft_array, axis=axes_to_reduce, keepdims=keepdims)
|
||||
|
||||
if features[key]["dtype"] in ["image", "video"]:
|
||||
value_norm = 1.0 if "depth" in key else 255.0
|
||||
ep_stats[key] = {
|
||||
k: v if k == "count" else np.squeeze(v / value_norm, axis=0) for k, v in ep_stats[key].items()
|
||||
}
|
||||
|
||||
return ep_stats
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user