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.reconstruction.plane import contact_table_height | |
| from real2sim.tracking import load_tracks | |
| def main(): | |
| 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, required=True) | |
| parser.add_argument("--table-height", type=float) | |
| parser.add_argument("--object-height", type=float, default=0.045) | |
| parser.add_argument("--contact-start", type=int, default=42) | |
| parser.add_argument("--contact-end", type=int, default=51) | |
| parser.add_argument("--contact-hand", choices=["left", "right"], default="right") | |
| args = parser.parse_args() | |
| folder = args.folder | |
| folder.mkdir(parents=True, exist_ok=True) | |
| episode = Episode.load(args.task, args.index) | |
| width, height = 518, 294 | |
| frames = video.read(episode.video_path) | |
| if len(frames) != len(episode): | |
| raise ValueError("Camera and video lengths must match") | |
| source_height, source_width = frames.shape[1:3] | |
| resized = np.stack( | |
| [ | |
| cv2.resize(frame, (width, height), interpolation=cv2.INTER_AREA) | |
| for frame in frames | |
| ] | |
| ) | |
| intrinsics = episode.intrinsics.copy() | |
| intrinsics[0] *= width / source_width | |
| intrinsics[1] *= height / source_height | |
| table_height = args.table_height | |
| if table_height is None: | |
| table_height = contact_table_height( | |
| episode.hands[args.contact_hand].keypoints, | |
| args.contact_start, | |
| args.contact_end, | |
| args.object_height, | |
| ) | |
| hands = { | |
| name: track.to_hand(episode.camera) | |
| for name, track in load_tracks(args.tracks).items() | |
| } | |
| scene = HandScene(hands, episode.camera, intrinsics, (width, height)) | |
| spec = scene.spec | |
| xml_path = folder / "scene.xml" | |
| floor = spec.geom("floor") | |
| floor.pos = [0, 0, table_height] | |
| floor.material = "" | |
| floor.rgba = [0.78, 0.74, 0.67, 1] | |
| model = spec.compile() | |
| xml_path.write_text(spec.to_xml()) | |
| mujoco.mj_saveModel(model, str(folder / "scene.mjb")) | |
| indices = np.arange(len(episode)) | |
| timestamps = indices / episode.fps | |
| video.write(folder / "real_video.mp4", resized, episode.fps) | |
| np.savez_compressed( | |
| folder / "sim_trajectory.npz", | |
| qpos=scene.qpos, | |
| camera=episode.camera, | |
| frame_indices=indices, | |
| timestamps=timestamps, | |
| ) | |
| np.savez_compressed( | |
| folder / "cameras.npz", | |
| **{ | |
| "camera/K": intrinsics, | |
| "camera/T_world_cam": episode.camera, | |
| "frame_indices": indices, | |
| "timestamps": timestamps, | |
| "image_size": np.array([width, height]), | |
| }, | |
| ) | |
| config = { | |
| "geometry_source": "flat_table", | |
| "camera_source": "egodex", | |
| "hand_source": "hawor", | |
| "task": args.task, | |
| "index": args.index, | |
| "frames": len(episode), | |
| "fps": episode.fps, | |
| "table_height_m": table_height, | |
| "object_height_prior_m": args.object_height, | |
| "table_source": "explicit" | |
| if args.table_height is not None | |
| else "egodex_hand_contact_prior", | |
| "contact_interval": [args.contact_start, args.contact_end], | |
| "contact_hand": args.contact_hand, | |
| "image_size": [width, height], | |
| "intrinsics": intrinsics.tolist(), | |
| } | |
| (folder / "scene.json").write_text(json.dumps(config, indent=2) + "\n") | |
| cv2.imwrite( | |
| str(folder / "real_preview.png"), | |
| cv2.cvtColor(resized[len(resized) // 2], cv2.COLOR_RGB2BGR), | |
| ) | |
| print(json.dumps(config), flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 3.89 kB
- Xet hash:
- 5caec31d05069b3b2324484f539d2f72385c6caf4b1aaa9c9b7e8bf93357ef64
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.