Refactor SaveScreenshot into a more useful RenderToTexture function.

Removes a few unused functions from helpers library.

PiperOrigin-RevId: 846661437
Change-Id: I815ffe1990098dedb71f694af088d85294f9c7de
This commit is contained in:
Haroon Qureshi
2025-12-19 03:33:07 -08:00
committed by Copybara-Service
parent 522ff1cd80
commit d1fd11bccd
5 changed files with 40 additions and 64 deletions
+4 -34
View File
@@ -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);