Add screenshot button to simulate.

PiperOrigin-RevId: 465628684
Change-Id: I0eb86d6b18c25bb437f197229be6d1c57ccbc50a

doc/changelog.rst

@@ -19,6 +19,7 @@ General
 - Added ``azimuth`` and ``elevation`` attributes to :ref:`visual/global<global>`, defining the initial orientation of
   the free camera at model load time.
  - Added ``mjv_defaultFreeCamera`` which sets the default free camera, respecting the above attributes.
+ - ``simulate`` now supports taking a screenshot via a button in the File section or via ``Ctrl-P``.
 
 Bug fixes
 ^^^^^^^^^
@@ -209,7 +210,7 @@ Version 2.1.3 (Mar. 23, 2022)
 General
 ^^^^^^^
 
-1. ``simulate`` now support cycling through cameras (with ``[`` and ``]`` keys).
+1. ``simulate`` now supports cycling through cameras (with the ``[`` and ``]`` keys).
 #. ``mjVIS_STATIC`` toggles all static bodies, not just direct children of the world.
 
 Python bindings

simulate/CMakeLists.txt

@@ -84,7 +84,7 @@ target_sources(
 target_include_directories(mjsimulate PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 target_compile_definitions(mjsimulate PUBLIC MJSIMULATE_STATIC)
 target_compile_options(mjsimulate PUBLIC ${MUJOCO_SIMULATE_COMPILE_OPTIONS})
-target_link_libraries(mjsimulate PUBLIC glfw mujoco::mujoco)
+target_link_libraries(mjsimulate PUBLIC glfw lodepng mujoco::mujoco)
 target_link_options(mjsimulate PRIVATE ${MUJOCO_SIMULATE_LINK_OPTIONS})
 
 if(APPLE)
@@ -120,6 +120,7 @@ target_link_libraries(
   mujoco::mujoco
   glfw
   Threads::Threads
+  lodepng
 )
 
 target_link_options(simulate PRIVATE ${MUJOCO_SIMULATE_LINK_OPTIONS})

simulate/simulate.cc

@@ -18,11 +18,11 @@
 #include <chrono>
 #include <cstdio>
 #include <cstring>
+#include <memory>
 #include <mutex>
 #include <string>
 #include <thread>
 
-
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mjxmacro.h>
@@ -78,6 +78,7 @@ const mjuiDef defFile[] = {
   {mjITEM_BUTTON,    "Print model",   2, nullptr,                    "CM"},
   {mjITEM_BUTTON,    "Print data",    2, nullptr,                    "CD"},
   {mjITEM_BUTTON,    "Quit",          1, nullptr,                    "CQ"},
+  {mjITEM_BUTTON,    "Screenshot",    2, NULL,                       "CP"},
   {mjITEM_END}
 };
 
@@ -1091,6 +1092,10 @@ void uiEvent(mjuiState* state) {
       case 4:             // Quit
         sim->exitrequest.store(1);
         break;
+
+      case 5:             // Screenshot
+        sim->screenshotrequest.store(true);
+        break;
       }
     }
 
@@ -1782,6 +1787,38 @@ void Simulate::render() {
     sensorshow(this, smallrect);
   }
 
+  // take screenshot, save to file
+  if (this->screenshotrequest.exchange(false)) {
+    const unsigned int h = uistate.rect[0].height;
+    const unsigned int w = uistate.rect[0].width;
+    std::unique_ptr<unsigned char[]> rgb(new unsigned char[3*w*h]);
+    if (!rgb) {
+      mju_error("could not allocate buffer for screenshot");
+    }
+    mjr_readPixels(rgb.get(), NULL, uistate.rect[0], &con);
+
+    // flip up-down
+    for (int r = 0; r < h/2; ++r) {
+      unsigned char* top_row = &rgb[3*w*r];
+      unsigned char* bottom_row = &rgb[3*w*(h-1-r)];
+      std::swap_ranges(top_row, top_row+3*w, bottom_row);
+    }
+
+    // save as PNG
+    // TODO(b/241577466): Parse the stem of the filename and use a .PNG extension.
+    // Unfortunately, if we just yank ".xml"/".mjb" from the filename and append .PNG, the macOS
+    // file dialog does not automatically open that location. Thus, we defer to a default
+    // "screenshot.png" for now.
+    const std::string path = getSavePath("screenshot.png");
+    if (!path.empty()) {
+      if (lodepng::encode(path, rgb.get(), w, h, LCT_RGB)) {
+        mju_error("could not save screenshot");
+      } else {
+        std::printf("saved screenshot: %s\n", path.c_str());
+      }
+    }
+  }
+
   // finalize
   glfwSwapBuffers(this->window);
 }

simulate/simulate.h

@@ -20,6 +20,7 @@
 #include <mutex>
 #include <thread>
 
+#include "lodepng.h"
 #include <GLFW/glfw3.h>
 #include <mujoco/mujoco.h>
 
@@ -112,6 +113,7 @@ class MJSIMULATEAPI Simulate {
   // 1: showing "loading" label, about to load
   // 0: model loaded or no load requested.
   int loadrequest = 0;
+  std::atomic_bool screenshotrequest = false;
   // strings
   char loadError[kMaxFilenameLength] = "";
   char dropfilename[kMaxFilenameLength] = "";