Use Ctrl+S to save camera configuration

Author: rohanpsingh

mujoco_viewer/callbacks.py

@@ -2,6 +2,7 @@
 import glfw
 import numpy as np
 import imageio
+import yaml
 from threading import Lock
 
 
@@ -54,7 +55,7 @@ def _key_callback(self, window, key, scancode, action, mods):
             self._advance_by_one_step = True
             self._paused = True
         # Slows down simulation
-        elif key == glfw.KEY_S:
+        elif key == glfw.KEY_S and mods != glfw.MOD_CONTROL:
             self._run_speed /= 2.0
         # Speeds up simulation
         elif key == glfw.KEY_F:
@@ -121,6 +122,22 @@ def _key_callback(self, window, key, scancode, action, mods):
         # Geom group visibility
         elif key in (glfw.KEY_0, glfw.KEY_1, glfw.KEY_2, glfw.KEY_3, glfw.KEY_4, glfw.KEY_5):
             self.vopt.geomgroup[key - glfw.KEY_0] ^= 1
+        elif key == glfw.KEY_S and mods == glfw.MOD_CONTROL:
+            cam_config = {
+                "type": self.cam.type,
+                "fixedcamid": self.cam.fixedcamid,
+                "trackbodyid": self.cam.trackbodyid,
+                "lookat": self.cam.lookat.tolist(),
+                "distance": self.cam.distance,
+                "azimuth": self.cam.azimuth,
+                "elevation": self.cam.elevation
+            }
+            try:
+                with open(self.CONFIG_PATH, "w") as f:
+                    yaml.dump(cam_config, f)
+                print("Camera config saved at {}".format(self.CONFIG_PATH))
+            except Exception as e:
+                print(e)
         # Quit
         if key == glfw.KEY_ESCAPE:
             print("Pressed ESC")

mujoco_viewer/mujoco_viewer.py

@@ -2,6 +2,8 @@
 import glfw
 import numpy as np
 import time
+import pathlib
+import yaml
 from .callbacks import Callbacks
 
 
@@ -27,6 +29,9 @@ def __init__(
         # keep true while running
         self.is_alive = True
 
+        self.CONFIG_PATH = pathlib.Path.joinpath(
+            pathlib.Path.home(), ".config/mujoco_viewer/config.yaml")
+
         # glfw init
         glfw.init()
 
@@ -68,6 +73,43 @@ def __init__(
         self.ctx = mujoco.MjrContext(
             self.model, mujoco.mjtFontScale.mjFONTSCALE_150.value)
 
+        # load camera from configuration (if available)
+        pathlib.Path(
+            self.CONFIG_PATH.parent).mkdir(
+            parents=True,
+            exist_ok=True)
+        pathlib.Path(self.CONFIG_PATH).touch(exist_ok=True)
+        with open(self.CONFIG_PATH, "r") as f:
+            try:
+                cam_config = {
+                    "type": self.cam.type,
+                    "fixedcamid": self.cam.fixedcamid,
+                    "trackbodyid": self.cam.trackbodyid,
+                    "lookat": self.cam.lookat.tolist(),
+                    "distance": self.cam.distance,
+                    "azimuth": self.cam.azimuth,
+                    "elevation": self.cam.elevation
+                }
+                load_config = yaml.safe_load(f)
+                if isinstance(load_config, dict):
+                    for key, val in load_config.items():
+                        if key in cam_config.keys():
+                            cam_config[key] = val
+                if cam_config["type"] == mujoco.mjtCamera.mjCAMERA_FIXED:
+                    if cam_config["fixedcamid"] < self.model.ncam:
+                        self.cam.type = cam_config["type"]
+                        self.cam.fixedcamid = cam_config["fixedcamid"]
+                if cam_config["type"] == mujoco.mjtCamera.mjCAMERA_TRACKING:
+                    if cam_config["trackbodyid"] < self.model.nbody:
+                        self.cam.type = cam_config["type"]
+                        self.cam.trackbodyid = cam_config["trackbodyid"]
+                self.cam.lookat = np.array(cam_config["lookat"])
+                self.cam.distance = cam_config["distance"]
+                self.cam.azimuth = cam_config["azimuth"]
+                self.cam.elevation = cam_config["elevation"]
+            except yaml.YAMLError as e:
+                print(e)
+
         # get viewport
         self.viewport = mujoco.MjrRect(
             0, 0, framebuffer_width, framebuffer_height)