From d1fd11bccd984e8346b78e955bbfb0587baf058a Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Fri, 19 Dec 2025 03:33:07 -0800 Subject: [PATCH] Refactor SaveScreenshot into a more useful RenderToTexture function. Removes a few unused functions from helpers library. PiperOrigin-RevId: 846661437 Change-Id: I815ffe1990098dedb71f694af088d85294f9c7de --- src/experimental/platform/helpers.cc | 38 +++------------------------ src/experimental/platform/helpers.h | 15 +++-------- src/experimental/platform/renderer.cc | 33 +++++++++++++++-------- src/experimental/platform/renderer.h | 10 +++---- src/experimental/studio/app.cc | 8 ++++-- 5 files changed, 40 insertions(+), 64 deletions(-) diff --git a/src/experimental/platform/helpers.cc b/src/experimental/platform/helpers.cc index a24c2d14..8b8e0ede 100644 --- a/src/experimental/platform/helpers.cc +++ b/src/experimental/platform/helpers.cc @@ -23,11 +23,9 @@ #include #include #include -#include #include "webp/encode.h" #include "webp/types.h" -#include #include #include #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(data), width, height, width * 3, &webp); std::ofstream file(filename, std::ios::binary); file.write(reinterpret_cast(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 byte_buffer; - byte_buffer.reserve(size * 3); - for (int i = 0; i < size; ++i) { - auto byte = static_cast(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(3 * width * height); - auto depth_buffer = std::vector(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); diff --git a/src/experimental/platform/helpers.h b/src/experimental/platform/helpers.h index f60bca0a..e880fb64 100644 --- a/src/experimental/platform/helpers.h +++ b/src/experimental/platform/helpers.h @@ -22,7 +22,6 @@ #include #include -#include #include namespace mujoco::platform { @@ -39,17 +38,9 @@ std::string LoadText(const std::string& filename); std::string ResolveFile(const std::string& filename, const std::vector& 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. diff --git a/src/experimental/platform/renderer.cc b/src/experimental/platform/renderer.cc index 9ab2a79d..15de97cd 100644 --- a/src/experimental/platform/renderer.cc +++ b/src/experimental/platform/renderer.cc @@ -14,10 +14,8 @@ #include "experimental/platform/renderer.h" -#include -#include +#include -#include "experimental/platform/helpers.h" #include 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 diff --git a/src/experimental/platform/renderer.h b/src/experimental/platform/renderer.h index d29f57ca..b89413bc 100644 --- a/src/experimental/platform/renderer.h +++ b/src/experimental/platform/renderer.h @@ -15,10 +15,8 @@ #ifndef MUJOCO_SRC_EXPERIMENTAL_PLATFORM_RENDERER_H_ #define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_RENDERER_H_ -#include -#include +#include #include -#include #include @@ -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; } diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index 3c61ae0e..5645ab64 100644 --- a/src/experimental/studio/app.cc +++ b/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 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();