Buckets:
| import mujoco | |
| import numpy as np | |
| from ..geometry import MUJOCO_CAMERA_FROM_OPENCV, look_at_xyaxes, quat_wxyz | |
| from .rig import SKIN, HandRig | |
| FLOOR_RGB = ([0.30, 0.32, 0.36], [0.40, 0.42, 0.46]) | |
| MAX_SIZE = (1920, 1080) | |
| class HandScene: | |
| def __init__(self, hands, camera, intrinsics, image_size, colors=None): | |
| self.hands = hands | |
| self.camera = camera | |
| colors = colors or {} | |
| self.rigs = [ | |
| HandRig(name, hand, colors.get(name, SKIN)) for name, hand in hands.items() | |
| ] | |
| self.fovy = float( | |
| np.degrees(2 * np.arctan(image_size[1] / 2 / intrinsics[1, 1])) | |
| ) | |
| self.model = self._compile() | |
| self.data = mujoco.MjData(self.model) | |
| self.qpos = sum(rig.qpos(self.model) for rig in self.rigs) | |
| self.mocap_pos = camera[:, :3, 3] | |
| self.mocap_quat = quat_wxyz(camera[:, :3, :3] @ MUJOCO_CAMERA_FROM_OPENCV) | |
| self.floor_id = self.model.geom("floor").id | |
| def from_episode(cls, episode, extra=None, colors=None, image_size=(1920, 1080)): | |
| hands = {f"egodex_{side}": hand for side, hand in episode.hands.items()} | |
| hands.update(extra or {}) | |
| return cls(hands, episode.camera, episode.intrinsics, image_size, colors) | |
| def __len__(self): | |
| return len(self.qpos) | |
| def _compile(self): | |
| spec = mujoco.MjSpec() | |
| spec.visual.global_.offwidth, spec.visual.global_.offheight = MAX_SIZE | |
| spec.visual.headlight.ambient = [0.45, 0.45, 0.45] | |
| spec.visual.headlight.diffuse = [0.6, 0.6, 0.6] | |
| spec.add_texture( | |
| name="grid", | |
| type=mujoco.mjtTexture.mjTEXTURE_2D, | |
| builtin=mujoco.mjtBuiltin.mjBUILTIN_CHECKER, | |
| rgb1=FLOOR_RGB[0], | |
| rgb2=FLOOR_RGB[1], | |
| width=256, | |
| height=256, | |
| ) | |
| material = spec.add_material(name="grid", texrepeat=[20, 20]) | |
| material.textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] = "grid" | |
| spec.worldbody.add_geom( | |
| name="floor", | |
| type=mujoco.mjtGeom.mjGEOM_PLANE, | |
| size=[10, 10, 0.1], | |
| material="grid", | |
| ) | |
| spec.worldbody.add_light(pos=[0, 0, 3], dir=[0, 0, -1], diffuse=[0.5, 0.5, 0.5]) | |
| for rig in self.rigs: | |
| rig.attach(spec) | |
| mount = spec.worldbody.add_body(name="ego_mount", mocap=True) | |
| mount.add_camera(name="ego", fovy=self.fovy) | |
| eye, target = self._overview_view() | |
| spec.worldbody.add_camera( | |
| name="overview", pos=eye, xyaxes=look_at_xyaxes(eye, target), fovy=45.0 | |
| ) | |
| return spec.compile() | |
| def _overview_view(self): | |
| wrists = np.concatenate([hand.positions[:, 0] for hand in self.hands.values()]) | |
| target = wrists.mean(axis=0) | |
| forward = self.camera[0, :3, 2].copy() | |
| forward[2] = 0.0 | |
| forward /= np.linalg.norm(forward) | |
| side = np.cross(forward, [0.0, 0.0, 1.0]) | |
| return target + 0.7 * side + 0.35 * np.array( | |
| [0.0, 0.0, 1.0] | |
| ) - 0.2 * forward, target | |
| def set_frame(self, frame): | |
| self.data.qpos[:] = self.qpos[frame] | |
| self.data.mocap_pos[0] = self.mocap_pos[frame] | |
| self.data.mocap_quat[0] = self.mocap_quat[frame] | |
| mujoco.mj_forward(self.model, self.data) | |
| def render(self, camera="ego", size=(960, 540)): | |
| renderer = mujoco.Renderer(self.model, height=size[1], width=size[0]) | |
| frames = [] | |
| for frame in range(len(self)): | |
| self.set_frame(frame) | |
| renderer.update_scene(self.data, camera=camera) | |
| frames.append(renderer.render().copy()) | |
| renderer.close() | |
| return np.stack(frames) | |
| def render_mask(self, camera="ego", size=(960, 540)): | |
| renderer = mujoco.Renderer(self.model, height=size[1], width=size[0]) | |
| renderer.enable_segmentation_rendering() | |
| masks = [] | |
| for frame in range(len(self)): | |
| self.set_frame(frame) | |
| renderer.update_scene(self.data, camera=camera) | |
| segmentation = renderer.render() | |
| is_geom = segmentation[..., 1] == mujoco.mjtObj.mjOBJ_GEOM | |
| masks.append(is_geom & (segmentation[..., 0] != self.floor_id)) | |
| renderer.close() | |
| return np.stack(masks).astype(np.float32) | |
| def joint_error(self): | |
| errors = {} | |
| for rig in self.rigs: | |
| size = rig.skeleton.size | |
| errors[rig.name] = np.full((len(self), size), np.nan) | |
| for frame in range(len(self)): | |
| self.set_frame(frame) | |
| for rig in self.rigs: | |
| simulated = np.stack( | |
| [ | |
| self.data.body(rig.body_name(j)).xpos | |
| for j in range(rig.skeleton.size) | |
| ] | |
| ) | |
| if rig.hand.visible[frame]: | |
| errors[rig.name][frame] = np.linalg.norm( | |
| simulated - rig.hand.positions[frame], axis=-1 | |
| ) | |
| return errors | |
Xet Storage Details
- Size:
- 5.04 kB
- Xet hash:
- f15e8315a852678e8df46f038c7403f379589595748da01d544dd9d1e1938499
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.