From 4271114ccad44a99b90bdf7cb55b3e07c3778881 Mon Sep 17 00:00:00 2001
From: Vikash Kumar <vikashplus@gmail.com>
Date: Thu, 24 Oct 2024 15:36:24 -0400
Subject: [PATCH] UPDATE:Upgrading the render_cams tutorial to official mujoco
 binding

---
 robohive/tutorials/render_cams.py | 42 +++++++++++++++----------------
 1 file changed, 20 insertions(+), 22 deletions(-)

diff --git a/robohive/tutorials/render_cams.py b/robohive/tutorials/render_cams.py
index c5ef5df9..551978d9 100644
--- a/robohive/tutorials/render_cams.py
+++ b/robohive/tutorials/render_cams.py
@@ -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'))
@@ -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")