Refactor SaveScreenshot into a more useful RenderToTexture function.

Removes a few unused functions from helpers library.

PiperOrigin-RevId: 846661437
Change-Id: I815ffe1990098dedb71f694af088d85294f9c7de

Author: haroonq

src/experimental/platform/helpers.cc

@@ -23,11 +23,9 @@
 #include <ios>
 #include <iterator>
 #include <string>
-#include <vector>
 
 #include "webp/encode.h"
 #include "webp/types.h"
-#include <mujoco/mjrender.h>
 #include <mujoco/mjxmacro.h>
 #include <mujoco/mujoco.h>
 #include "engine/engine_vis_visualize.h"
@@ -89,45 +87,17 @@ std::string ResolveFile(const std::string& filename,
   return "";
 }
 
-void SaveColorToWebp(int width, int height, const unsigned char* data,
-                     const std::string& filename) {
+void SaveToWebp(int width, int height, const std::byte* data,
+                const std::string& filename) {
   uint8_t* webp = nullptr;
-  const size_t size =
-      WebPEncodeLosslessRGB(data, width, height, width * 3, &webp);
-
+  const size_t size = WebPEncodeLosslessRGB(
+      reinterpret_cast<const uint8_t*>(data), width, height, width * 3, &webp);
   std::ofstream file(filename, std::ios::binary);
   file.write(reinterpret_cast<const char*>(webp), size);
   file.close();
   WebPFree(webp);
 }
 
-void SaveDepthToWebp(int width, int height, const float* data,
-                     const std::string& filename) {
-  const int size = width * height;
-
-  // Turn the depth buffer into a greyscale color buffer.
-  std::vector<unsigned char> byte_buffer;
-  byte_buffer.reserve(size * 3);
-  for (int i = 0; i < size; ++i) {
-    auto byte = static_cast<int>(255.0 * data[i]);
-    byte_buffer.push_back(byte);
-    byte_buffer.push_back(byte);
-    byte_buffer.push_back(byte);
-  }
-  SaveColorToWebp(width, height, byte_buffer.data(), filename);
-}
-
-void SaveScreenshotToWebp(int width, int height, mjrContext* con,
-                          const std::string& filename) {
-  mjr_setBuffer(mjFB_OFFSCREEN, con);
-  auto rgb_buffer = std::vector<unsigned char>(3 * width * height);
-  auto depth_buffer = std::vector<float>(width * height, 1.0f);
-  mjrRect viewport = {0, 0, width, height};
-  mjr_readPixels(rgb_buffer.data(), depth_buffer.data(), viewport, con);
-  mjr_setBuffer(mjFB_WINDOW, con);
-  SaveColorToWebp(width, height, rgb_buffer.data(), filename);
-}
-
 const void* GetValue(const mjModel* model, const mjData* data,
                      const char* field, int index) {
   MJDATA_POINTERS_PREAMBLE(model);

src/experimental/platform/helpers.h

@@ -22,7 +22,6 @@
 #include <string_view>
 #include <vector>
 
-#include <mujoco/mjrender.h>
 #include <mujoco/mujoco.h>
 
 namespace mujoco::platform {
@@ -39,17 +38,9 @@ std::string LoadText(const std::string& filename);
 std::string ResolveFile(const std::string& filename,
                         const std::vector<std::string>& search_paths);
 
-// Exports the given color buffer to a webp file.
-void SaveColorToWebp(int width, int height, const unsigned char* data,
-                     const std::string& filename);
-
-// Exports the given depth buffer to a webp file.
-void SaveDepthToWebp(int width, int height, const float* data,
-                     const std::string& filename);
-
-// Exports the current state of the mjrContext to a webp file.
-void SaveScreenshotToWebp(int width, int height, mjrContext* con,
-                          const std::string& filename);
+// Exports the given image (assumed to be RGB888) to a webp file.
+void SaveToWebp(int width, int height, const std::byte* data,
+                const std::string& filename);
 
 // Returns a pointer to the value of the given field in the given data.
 // Returns nullptr if the field is not found or the index is out of bounds.

src/experimental/platform/renderer.cc

@@ -14,10 +14,8 @@
 
 #include "experimental/platform/renderer.h"
 
-#include <chrono>
-#include <string>
+#include <cstddef>
 
-#include "experimental/platform/helpers.h"
 #include <mujoco/mujoco.h>
 
 namespace mujoco::platform {
@@ -50,18 +48,31 @@ void Renderer::Deinit() {
 void Renderer::Render(const mjModel* model, mjData* data,
                       const mjvPerturb* perturb, mjvCamera* camera,
                       const mjvOption* vis_option, int width, int height) {
-  if (initialized_) {
-    mjv_updateScene(model, data, vis_option, perturb, camera, mjCAT_ALL,
-                    &scene_);
+  if (!initialized_) {
+    return;
   }
 
-  mjrRect main_viewport = {0, 0, width, height};
-  mjr_render(main_viewport, data ? &scene_ : nullptr, &render_context_);
+  mjv_updateScene(model, data, vis_option, perturb, camera, mjCAT_ALL,
+                  &scene_);
+
+  const mjrRect viewport = {0, 0, width, height};
+  mjr_render(viewport, &scene_, &render_context_);
 }
 
-void Renderer::SaveScreenshot(const std::string& filename, int width,
-                              int height) {
-  SaveScreenshotToWebp(width, height, &render_context_, filename);
+void Renderer::RenderToTexture(const mjModel* model, mjData* data,
+                               mjvCamera* camera, int width, int height,
+                               std::byte* output) {
+  if (!initialized_) {
+    return;
+  }
+
+  const mjrRect viewport = {0, 0, width, height};
+
+  mjr_setBuffer(mjFB_OFFSCREEN, &render_context_);
+  mjv_updateCamera(model, data, camera, &scene_);
+  mjr_render(viewport, &scene_, &render_context_);
+  mjr_readPixels((unsigned char*)output, nullptr, viewport, &render_context_);
+  mjr_setBuffer(mjFB_WINDOW, &render_context_);
 }
 
 }  // namespace mujoco::platform

src/experimental/platform/renderer.h

@@ -15,10 +15,8 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_PLATFORM_RENDERER_H_
 #define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_RENDERER_H_
 
-#include <chrono>
-#include <cstdint>
+#include <cstddef>
 #include <functional>
-#include <string>
 
 #include <mujoco/mujoco.h>
 
@@ -48,8 +46,10 @@ class Renderer {
               mjvCamera* camera, const mjvOption* vis_option, int width,
               int height);
 
-  // Saves a screenshot of the simulation state into the given file.
-  void SaveScreenshot(const std::string& filename, int width, int height);
+  // Populates the given output buffer with RGB888 pixel data. The size of the
+  // output buffer must be at least width * height * 3.
+  void RenderToTexture(const mjModel* model, mjData* data, mjvCamera* camera,
+                       int width, int height, std::byte* output);
 
   // Rendering flags.
   mjtByte* GetRenderFlags() { return scene_.flags; }

src/experimental/studio/app.cc

@@ -1678,8 +1678,12 @@ void App::FileDialogGui() {
   if (ImGui::BeginPopupModal("SaveWebp", NULL,
                              ImGuiWindowFlags_AlwaysAutoResize)) {
     if (platform::ImGui_FileDialog(tmp_.filename, sizeof(tmp_.filename))) {
-      renderer_->SaveScreenshot(tmp_.filename, window_->GetWidth(),
-                                window_->GetHeight());
+      const int width = window_->GetWidth();
+      const int height = window_->GetHeight();
+      std::vector<std::byte> buffer(width * height * 3);
+      renderer_->RenderToTexture(model_, data_, &camera_, width, height,
+                                 buffer.data());
+      platform::SaveToWebp(width, height, buffer.data(), tmp_.filename);
       tmp_.last_save_screenshot_file = tmp_.filename;
     }
     ImGui::EndPopup();