added depth to read_pixels

Author: LuanGBR

mujoco_viewer/mujoco_viewer.py

@@ -268,7 +268,7 @@ def apply_perturbations(self):
         mujoco.mjv_applyPerturbPose(self.model, self.data, self.pert, 0)
         mujoco.mjv_applyPerturbForce(self.model, self.data, self.pert)
 
-    def read_pixels(self, camid=None):
+    def read_pixels(self, camid=None,depth=False):
         if self.render_mode == 'window':
             raise NotImplementedError(
                 "Use 'render()' in 'window' mode.")
@@ -293,13 +293,17 @@ def read_pixels(self, camid=None):
             self.scn)
         # render
         mujoco.mjr_render(self.viewport, self.scn, self.ctx)
+        shape = glfw.get_framebuffer_size(self.window)
+        img = np.zeros((shape[1], shape[0], 3), dtype=np.uint8)
+        
+        if depth:
+            depth_img = np.zeros((shape[1], shape[0], 1), dtype=np.uint8)
+            mujoco.mjr_readPixels(img, depth_img, self.viewport, self.ctx)
+            return np.dstack((np.flipud(img),np.flipud(depth_img)))
+        else:
+            mujoco.mjr_readPixels(img, None, self.viewport, self.ctx)
+            return np.flipud(img)
 
-        img = np.zeros(
-            (glfw.get_framebuffer_size(
-                self.window)[1], glfw.get_framebuffer_size(
-                self.window)[0], 3), dtype=np.uint8)
-        mujoco.mjr_readPixels(img, None, self.viewport, self.ctx)
-        return np.flipud(img)
 
     def render(self):
         if self.render_mode == 'offscreen':