Buckets:
| import argparse | |
| from pathlib import Path | |
| import cv2 | |
| import mujoco | |
| import numpy as np | |
| from real2sim import video | |
| from real2sim.geometry import MUJOCO_CAMERA_FROM_OPENCV, quat_wxyz | |
| def render(folder, camera, size): | |
| model = mujoco.MjModel.from_binary_path(str(folder / "scene.mjb")) | |
| data = mujoco.MjData(model) | |
| with np.load(folder / "rollout.npz") as saved: | |
| qpos = saved["qpos"] | |
| mount = int(model.body_mocapid[model.body("ego_mount").id]) | |
| frames = [] | |
| with mujoco.Renderer(model, height=size[1], width=size[0]) as renderer: | |
| for index, pose in enumerate(qpos): | |
| data.qpos[:] = pose | |
| data.mocap_pos[mount] = camera[index, :3, 3] | |
| data.mocap_quat[mount] = quat_wxyz( | |
| camera[index, :3, :3] @ MUJOCO_CAMERA_FROM_OPENCV | |
| ) | |
| mujoco.mj_forward(model, data) | |
| renderer.update_scene(data, camera="ego") | |
| frames.append(renderer.render().copy()) | |
| return np.asarray(frames) | |
| def label(frame, text): | |
| image = frame.copy() | |
| cv2.rectangle(image, (0, 0), (image.shape[1], 28), (18, 22, 30), -1) | |
| cv2.putText( | |
| image, | |
| text, | |
| (10, 20), | |
| cv2.FONT_HERSHEY_SIMPLEX, | |
| 0.55, | |
| (245, 245, 245), | |
| 1, | |
| cv2.LINE_AA, | |
| ) | |
| return image | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("folder", type=Path) | |
| args = parser.parse_args() | |
| root = args.folder | |
| source = video.read(root / "mujoco_replay.mp4") | |
| size = (source.shape[2] // 3, source.shape[1]) | |
| source = source[:, :, : size[0]] | |
| with np.load(root / "sim_trajectory.npz") as saved: | |
| camera, timestamps = saved["camera"], saved["timestamps"] | |
| initial = render(root / "physics/initial", camera, size) | |
| best = render(root / "physics/refinement", camera, size) | |
| frames = [ | |
| np.concatenate( | |
| ( | |
| label(real, "Real"), | |
| label(before, "Initial contact physics"), | |
| label(after, "Pixel-guided contact physics"), | |
| ), | |
| axis=1, | |
| ) | |
| for real, before, after in zip(source, initial, best) | |
| ] | |
| video.write( | |
| root / "final_comparison.mp4", frames, 1 / float(np.median(np.diff(timestamps))) | |
| ) | |
| cv2.imwrite( | |
| str(root / "final_preview.png"), | |
| cv2.cvtColor(frames[len(frames) // 2], cv2.COLOR_RGB2BGR), | |
| ) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 2.45 kB
- Xet hash:
- 562e383c3dee31b457ce184769f7e79d1fc2682e2fec298853d93ade383562e7
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.