Buckets:
| import argparse | |
| import json | |
| from pathlib import Path | |
| import cv2 | |
| import mujoco | |
| import numpy as np | |
| from real2sim import Episode, HandScene, video | |
| from real2sim.hands import Hand | |
| from real2sim.reconstruction.preview import load_reconstruction | |
| from real2sim.reconstruction.surface import ( | |
| anchor_transform, | |
| transform_points, | |
| triangulate, | |
| ) | |
| from real2sim.sim.reconstructed import ReconstructedScene | |
| def main(argv=None): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("folder", type=Path) | |
| parser.add_argument("--task", default="basic_pick_place") | |
| parser.add_argument("--index", type=int, default=0) | |
| parser.add_argument("--tracks", type=Path) | |
| args = parser.parse_args(argv) | |
| result = load_reconstruction(args.folder) | |
| episode = Episode.load(args.task, args.index) | |
| indices = result.frame_indices | |
| if len(indices) < 2 or indices[-1] >= len(episode): | |
| raise ValueError( | |
| "Reconstruction must contain at least two frames within the episode" | |
| ) | |
| view = result.views[0] | |
| replay_indices = np.arange(indices[0], indices[-1] + 1) | |
| timestamps = replay_indices / episode.fps | |
| height, width = view["mask"].shape | |
| image_size = (width, height) | |
| transform = anchor_transform(view["camera_poses"], episode.camera[indices[0]]) | |
| hands = { | |
| name: Hand( | |
| hand.side, | |
| hand.poses[replay_indices], | |
| hand.confidence[replay_indices], | |
| hand.skeleton, | |
| hand.visible[replay_indices], | |
| ) | |
| for name, hand in episode.hands.items() | |
| } | |
| if args.tracks: | |
| from real2sim.tracking import load_tracks | |
| tracked = { | |
| name: track.to_hand(episode.camera) | |
| for name, track in load_tracks(args.tracks).items() | |
| } | |
| hands = { | |
| name: Hand( | |
| hand.side, | |
| hand.poses[replay_indices], | |
| hand.confidence[replay_indices], | |
| hand.skeleton, | |
| hand.visible[replay_indices], | |
| ) | |
| for name, hand in tracked.items() | |
| } | |
| camera = episode.camera[replay_indices] | |
| initial = { | |
| name: Hand( | |
| hand.side, | |
| hand.poses[:1], | |
| hand.confidence[:1], | |
| hand.skeleton, | |
| hand.visible[:1], | |
| ) | |
| for name, hand in hands.items() | |
| } | |
| scene = HandScene(initial, camera[:1], view["intrinsics"], image_size) | |
| hand_mask = scene.render_mask(size=image_size)[0] > 0 | |
| excluded = cv2.dilate(hand_mask.astype(np.uint8), np.ones((9, 9), np.uint8)).astype( | |
| bool | |
| ) | |
| vertices, faces, uv = triangulate(view, exclude=excluded) | |
| vertices = transform_points(vertices, transform) | |
| scene = ReconstructedScene( | |
| hands, | |
| camera, | |
| view["intrinsics"], | |
| image_size, | |
| vertices, | |
| faces, | |
| uv, | |
| args.folder / "view_0000.png", | |
| ) | |
| ego, overview = scene.render(size=image_size), scene.render("overview", image_size) | |
| from mapanything.utils.image import preprocess_inputs, rgb | |
| frames = video.read(episode.video_path, limit=int(replay_indices[-1]) + 1) | |
| processed = preprocess_inputs([{"img": frames[i]} for i in replay_indices]) | |
| real = np.stack( | |
| [ | |
| (rgb(v["img"], "dinov2")[0] * 255).clip(0, 255).astype(np.uint8) | |
| for v in processed | |
| ] | |
| ) | |
| ids = scene.render_ids(size=image_size) | |
| hand_ids = [ | |
| i | |
| for i in range(scene.model.ngeom) | |
| if scene.model.geom(i).name not in ("floor", "scene_surface") | |
| ] | |
| hand_masks = np.isin(ids, hand_ids) | |
| video.write( | |
| args.folder / "mujoco_replay.mp4", | |
| video.hstack(real, ego, overview), | |
| episode.fps, | |
| ) | |
| middle = len(replay_indices) // 2 | |
| cv2.imwrite( | |
| str(args.folder / "mujoco_preview.png"), | |
| cv2.cvtColor( | |
| np.concatenate((real[middle], ego[middle], overview[middle]), axis=1), | |
| cv2.COLOR_RGB2BGR, | |
| ), | |
| ) | |
| np.savez_compressed( | |
| args.folder / "surface.npz", | |
| vertices=vertices, | |
| faces=faces, | |
| uv=uv, | |
| T_world_reconstruction=transform, | |
| excluded=excluded, | |
| ) | |
| np.savez_compressed( | |
| args.folder / "sim_masks.npz", | |
| geom_ids=ids, | |
| hand=hand_masks, | |
| frame_indices=replay_indices, | |
| timestamps=timestamps, | |
| geom_names=np.asarray( | |
| [scene.model.geom(i).name for i in range(scene.model.ngeom)] | |
| ), | |
| ) | |
| trajectory = { | |
| "qpos": scene.qpos, | |
| "camera": camera, | |
| "frame_indices": replay_indices, | |
| "timestamps": timestamps, | |
| } | |
| np.savez_compressed(args.folder / "sim_trajectory.npz", **trajectory) | |
| xml = scene.spec.to_xml().replace(scene.texture, "view_0000.png") | |
| (args.folder / "scene.xml").write_text(xml) | |
| mujoco.MjModel.from_xml_path(str(args.folder / "scene.xml")) | |
| mujoco.mj_saveModel(scene.model, str(args.folder / "scene.mjb")) | |
| reloaded = mujoco.MjModel.from_binary_path(str(args.folder / "scene.mjb")) | |
| reloaded_data = mujoco.MjData(reloaded) | |
| roundtrip_errors = [] | |
| with mujoco.Renderer(reloaded, height=height, width=width) as renderer: | |
| renderer.enable_segmentation_rendering() | |
| for index in (0, len(replay_indices) // 2, len(replay_indices) - 1): | |
| reloaded_data.qpos[:] = scene.qpos[index] | |
| reloaded_data.mocap_pos[0] = scene.mocap_pos[index] | |
| reloaded_data.mocap_quat[0] = scene.mocap_quat[index] | |
| mujoco.mj_forward(reloaded, reloaded_data) | |
| renderer.update_scene(reloaded_data, camera="ego") | |
| segmentation = renderer.render() | |
| rendered_ids = np.where( | |
| segmentation[..., 1] == mujoco.mjtObj.mjOBJ_GEOM, | |
| segmentation[..., 0], | |
| -1, | |
| ) | |
| roundtrip_errors.append(int(np.count_nonzero(rendered_ids != ids[index]))) | |
| if any(roundtrip_errors): | |
| raise RuntimeError(f"Scene reload changed renderer masks: {roundtrip_errors}") | |
| errors = scene.joint_error() | |
| predicted = np.stack([transform @ v["camera_poses"] for v in result.views]) | |
| metrics = { | |
| "vertices": len(vertices), | |
| "triangles": len(faces), | |
| "hand_source": "hawor" if args.tracks else "egodex", | |
| "simulation_mode": "kinematic_replay", | |
| "scene_surface": "static_visual_mesh", | |
| "scene_reload_mask_errors": roundtrip_errors, | |
| "replay_frames": len(replay_indices), | |
| "fps": episode.fps, | |
| "camera_translation_rmse_m": float( | |
| np.sqrt( | |
| np.mean( | |
| np.sum( | |
| (predicted[:, :3, 3] - episode.camera[indices, :3, 3]) ** 2, | |
| axis=1, | |
| ) | |
| ) | |
| ) | |
| ), | |
| "hand_joint_error_max_m": {k: float(np.nanmax(v)) for k, v in errors.items()}, | |
| } | |
| (args.folder / "sim_metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") | |
| print(json.dumps(metrics)) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 7.11 kB
- Xet hash:
- e41f531f417370da24ad7a883b6096a950b419d0a8af25790a65661d9e511a31
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.