diff --git a/README.md b/README.md index 0d47d24..aa15ca1 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,11 @@ img = viewer.read_pixels(camid=2) ## do something cool with img ``` +# Other features +1. You can save a screenshot of the viewer programmatically: + `viewer.save_screenshot('path/to/save.png')`. + + # Optional Parameters - `title`: set the title of the window, for example: `viewer = mujoco_viewer.MujocoViewer(model, data, title='My Demo')` (defaults to `mujoco-python-viewer`). diff --git a/mujoco_viewer/callbacks.py b/mujoco_viewer/callbacks.py index 92c63b8..621c4c6 100644 --- a/mujoco_viewer/callbacks.py +++ b/mujoco_viewer/callbacks.py @@ -34,9 +34,24 @@ def __init__(self, hide_menus): self._loop_count = 0 self._advance_by_one_step = False self._hide_menus = hide_menus + self._left_shift_pressed = False + self.user_option_pressed = None + + def save_screenshot(self, file_name): + 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) + imageio.imwrite(file_name, np.flipud(img)) def _key_callback(self, window, key, scancode, action, mods): - if action != glfw.RELEASE: + if key == glfw.KEY_LEFT_SHIFT: + self._left_shift_pressed = not self._left_shift_pressed + + if self._left_shift_pressed and key != glfw.KEY_LEFT_SHIFT: + self.user_option_pressed = chr(key) + elif action != glfw.RELEASE: if key == glfw.KEY_LEFT_ALT: self._hide_menus = False return @@ -65,12 +80,7 @@ def _key_callback(self, window, key, scancode, action, mods): self._render_every_frame = not self._render_every_frame # Capture screenshot elif key == glfw.KEY_T: - 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) - imageio.imwrite(self._image_path % self._image_idx, np.flipud(img)) + self.save_screenshot(self._image_path % self._image_idx) self._image_idx += 1 # Display contact forces elif key == glfw.KEY_C: diff --git a/mujoco_viewer/mujoco_viewer.py b/mujoco_viewer/mujoco_viewer.py index 75c7355..c490eeb 100644 --- a/mujoco_viewer/mujoco_viewer.py +++ b/mujoco_viewer/mujoco_viewer.py @@ -118,6 +118,9 @@ def __init__( self._overlay = {} self._markers = [] + def set_title(self, new_title): + glfw.set_window_title(self.window, new_title) + def add_marker(self, **marker_params): self._markers.append(marker_params)