chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Save depth images in visualize_render.py

Browse Source 
PiperOrigin-RevId: 871454618
Change-Id: If844e2c4af10e271536c4e6a36e760e8d4d50374
This commit is contained in:
Baruch Tabanpour
2026-02-17 12:54:44 -08:00
committed by Copybara-Service
parent 448e221690
commit f6129596a6
1 changed files with 18 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/warp/visualize_render.py
+18 -1
View File
@@ -175,23 +175,40 @@ def _main(_: Sequence[str]):
)(mx, dx_batch, rc) )(mx, dx_batch, rc)
rgb_packed = out_batch[0] rgb_packed = out_batch[0]
print(f' rgb shape: {rgb_packed.shape}\n') depth_packed = out_batch[1]
print(f' rgb shape: {rgb_packed.shape}')
print(f' depth shape: {depth_packed.shape}\n')
rgb = jax.vmap( rgb = jax.vmap(
render_util.get_rgb, in_axes=(None, 0, None) render_util.get_rgb, in_axes=(None, 0, None)
)(rc, rgb_packed, _CAMERA_ID.value) )(rc, rgb_packed, _CAMERA_ID.value)
depth = jax.vmap(
render_util.get_depth, in_axes=(None, 0, None, None)
)(rc, depth_packed, _CAMERA_ID.value, 10.0)
single_path = os.path.join( single_path = os.path.join(
_OUTPUT_DIR.value, f'camera_{_CAMERA_ID.value}.png' _OUTPUT_DIR.value, f'camera_{_CAMERA_ID.value}.png'
) )
_save_single(rgb, single_path) _save_single(rgb, single_path)
depth_rgb = np.repeat(np.asarray(depth)[..., None], 3, axis=-1)
depth_single_path = os.path.join(
_OUTPUT_DIR.value, f'depth_{_CAMERA_ID.value}.png'
)
_save_single(depth_rgb, depth_single_path)
if _NWORLD.value > 1: if _NWORLD.value > 1:
tiled_path = os.path.join( tiled_path = os.path.join(
_OUTPUT_DIR.value, f'tiled_{_CAMERA_ID.value}.png' _OUTPUT_DIR.value, f'tiled_{_CAMERA_ID.value}.png'
) )
_save_tiled(rgb, tiled_path) _save_tiled(rgb, tiled_path)
depth_tiled_path = os.path.join(
_OUTPUT_DIR.value, f'depth_tiled_{_CAMERA_ID.value}.png'
)
_save_tiled(depth_rgb, depth_tiled_path)
print('\ndone.') print('\ndone.')