Skip to content

Commit

Permalink
UPDATE:Upgrading the render_cams tutorial to official mujoco binding
Browse files Browse the repository at this point in the history
  • Loading branch information
vikashplus committed Oct 24, 2024
1 parent a9d227e commit 4271114
Showing 1 changed file with 20 additions and 22 deletions.
42 changes: 20 additions & 22 deletions robohive/tutorials/render_cams.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,29 +9,18 @@
DESC = '''
Helper script to render images offscreen and save using a mujoco model.\n
USAGE:\n
$ python render_cams.py --model_path franka_sim.xml --cam_names top_cam --cam_names left_cam --cam_names right_cam \n
$ python render_cams.py --model_path <../model.xml> --cam_names <name1> --cam_names <name2> \n
EXAMPLE:\n
$ python utils/render_cams.py -m envs/fm/assets/franka_microwave.xml -c top_cam -c left_cam -c right_cam
$ python robohive/tutorials/render_cams.py -m robohive/envs/arms/franka/assets/franka_reach_v0.xml -c left_cam -c top_cam
$ python robohive/tutorials/render_cams.py -m robohive/simhive/robel_sim/dkitty/kitty-v2.1.xml -c "A:trackingY" -c "A:trackingZ"
'''

from mujoco_py import load_model_from_path, MjSim
import mujoco

from PIL import Image
import numpy as np
import click


def render_camera_offscreen(cameras:list, width:int=640, height:int=480, device_id:int=0, sim=None):
"""
Render images(widthxheight) from a list_of_cameras on the specified device_id.
"""
imgs = np.zeros((len(cameras), height, width, 3), dtype=np.uint8)
for ind, cam in enumerate(cameras) :
img = sim.render(width=width, height=height, mode='offscreen', camera_name=cam, device_id=device_id)
img = img[::-1, :, : ] # Image has to be flipped
imgs[ind, :, :, :] = img
return imgs


@click.command(help=DESC)
@click.option('-m', '--model_path', required=True, type=str, help='model file')
@click.option('-c', '--cam_names', required=True, multiple=True, help=('Camera names for rendering'))
Expand All @@ -40,14 +29,23 @@ def render_camera_offscreen(cameras:list, width:int=640, height:int=480, device_
@click.option('-d', '--device_id', type=int, default=0, help='device id for rendering')

def main(model_path, cam_names, width, height, device_id):
# render images
model = load_model_from_path(model_path)
sim = MjSim(model)
imgs = render_camera_offscreen(cameras=cam_names, width=width, height=height, device_id=device_id, sim=sim)

# prepare model, data, scene
mj_model = mujoco.MjModel.from_xml_path(model_path)
mj_data = mujoco.MjData(mj_model)
mujoco.mj_forward(mj_model, mj_data)

# prepare the renderer
renderer = mujoco.Renderer(mj_model, height=height, width=width)

# save images
for i, cam in enumerate(cam_names):
image = Image.fromarray(imgs[i])
# update the scene
renderer.update_scene(mj_data, camera=cam)
# render the rgb_array
rgb_arr = renderer.render()
# save the image
image = Image.fromarray(rgb_arr)
image.save(cam+".jpeg")
print("saved "+cam+".jpeg")

Expand Down

0 comments on commit 4271114

Please sign in to comment.