chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allow flipping textures with any number of channels and user provided pixel data.

Browse Source 
PiperOrigin-RevId: 813256285
Change-Id: I0d49885dc483ab8f093b2566ec4bef355575505e
This commit is contained in:
Sam Haves
2025-09-30 07:27:12 -07:00
committed by Copybara-Service
parent bdccbcc60a
commit 2dc533b63b
3 changed files with 128 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/user/user_objects.cc
+67 -85
View File
@@ -20,6 +20,7 @@
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <functional>
@@ -59,8 +60,9 @@ class PNGImage {
int Height() const { return height_; }
bool IsSRGB() const { return is_srgb_; }
uint8_t operator[] (int i) const { return data_[i]; }
std::vector<unsigned char>& MoveData() { return data_; }
std::byte operator[] (int i) const { return data_[i]; }
mjByteVec&& MoveData() && { return std::move(data_); }
private:
std::size_t Size() const {
@@ -71,7 +73,7 @@ class PNGImage {
int height_;
bool is_srgb_;
LodePNGColorType color_type_;
std::vector<uint8_t> data_;
mjByteVec data_;
};
PNGImage PNGImage::Load(const mjCBase* obj, mjResource* resource,
@@ -113,7 +115,12 @@ PNGImage PNGImage::Load(const mjCBase* obj, mjResource* resource,
lodepng::State state;
state.info_raw.colortype = image.color_type_;
state.info_raw.bitdepth = 8;
unsigned err = lodepng::decode(image.data_, w, h, state, buffer, nbuffer);
unsigned char* data_ptr = nullptr;
unsigned err = lodepng_decode(&data_ptr, &w, &h, &state, buffer, nbuffer);
struct free_delete {
void operator()(unsigned char* ptr) const { std::free(ptr); }
};
std::unique_ptr<unsigned char, free_delete> data{data_ptr};
// check for errors
if (err) {
@@ -122,6 +129,13 @@ PNGImage PNGImage::Load(const mjCBase* obj, mjResource* resource,
throw mjCError(obj, "%s", ss.str().c_str());
}
if (data) {
size_t buffersize = lodepng_get_raw_size(w, h, &state.info_raw);
image.data_.insert(image.data_.end(),
reinterpret_cast<std::byte*>(data.get()),
reinterpret_cast<std::byte*>(&data.get()[buffersize]));
}
image.width_ = w;
image.height_ = h;
image.is_srgb_ = (state.info_png.srgb_defined == 1);
@@ -4890,7 +4904,7 @@ void mjCTexture::BuiltinCube(void) {
// load PNG file
void mjCTexture::LoadPNG(mjResource* resource,
std::vector<unsigned char>& image,
std::vector<std::byte>& image,
unsigned int& w, unsigned int& h, bool& is_srgb) {
LodePNGColorType color_type;
if (nchannel == 4) {
@@ -4907,13 +4921,14 @@ void mjCTexture::LoadPNG(mjResource* resource,
w = png_image.Width();
h = png_image.Height();
is_srgb = png_image.IsSRGB();
image = png_image.MoveData();
// Move data into image.
image = std::move(png_image).MoveData();
}
// load KTX file
void mjCTexture::LoadKTX(mjResource* resource,
std::vector<unsigned char>& image, unsigned int& w,
unsigned int& h, bool& is_srgb) {
void mjCTexture::LoadKTX(mjResource* resource, std::vector<std::byte>& image,
unsigned int& w, unsigned int& h, bool& is_srgb) {
const void* buffer = 0;
int buffer_sz = mju_readResource(resource, &buffer);
@@ -4933,8 +4948,7 @@ void mjCTexture::LoadKTX(mjResource* resource,
}
// load custom file
void mjCTexture::LoadCustom(mjResource* resource,
std::vector<unsigned char>& image,
void mjCTexture::LoadCustom(mjResource* resource, std::vector<std::byte>& image,
unsigned int& w, unsigned int& h, bool& is_srgb) {
const void* buffer = 0;
int buffer_sz = mju_readResource(resource, &buffer);
@@ -4972,11 +4986,44 @@ void mjCTexture::LoadCustom(mjResource* resource,
memcpy(image.data(), (void*)(pint+2), w*h*3*sizeof(char));
}
void mjCTexture::FlipIfNeeded(std::vector<std::byte>& image, unsigned int w,
unsigned int h) {
// horizontal flip
if (hflip) {
for (int r = 0; r < h; r++) {
for (int c = 0; c < w / 2; c++) {
int c1 = w - 1 - c;
auto val1 = nchannel * (r * w + c);
auto val2 = nchannel * (r * w + c1);
for (int ch = 0; ch < nchannel; ch++) {
auto tmp = image[val1 + ch];
image[val1 + ch] = image[val2 + ch];
image[val2 + ch] = tmp;
}
}
}
}
// vertical flip
if (vflip) {
for (int r = 0; r < h / 2; r++) {
for (int c = 0; c < w; c++) {
int r1 = h - 1 - r;
auto val1 = nchannel * (r * w + c);
auto val2 = nchannel * (r1 * w + c);
for (int ch = 0; ch < nchannel; ch++) {
auto tmp = image[val1 + ch];
image[val1 + ch] = image[val2 + ch];
image[val2 + ch] = tmp;
}
}
}
}
}
// load from PNG or custom file, flip if specified
void mjCTexture::LoadFlip(std::string filename, const mjVFS* vfs,
std::vector<unsigned char>& image,
std::vector<std::byte>& image,
unsigned int& w, unsigned int& h, bool& is_srgb) {
std::string asset_type = GetAssetContentType(filename, content_type_);
@@ -5008,68 +5055,16 @@ void mjCTexture::LoadFlip(std::string filename, const mjVFS* vfs,
throw err;
}
// horizontal flip
if (hflip) {
if (nchannel != 3) {
throw mjCError(
this, "currently only 3-channel textures support horizontal flip");
}
for (int r=0; r < h; r++) {
for (int c=0; c < w/2; c++) {
int c1 = w-1-c;
unsigned char tmp[3] = {
image[3*(r*w+c)],
image[3*(r*w+c)+1],
image[3*(r*w+c)+2]
};
image[3*(r*w+c)] = image[3*(r*w+c1)];
image[3*(r*w+c)+1] = image[3*(r*w+c1)+1];
image[3*(r*w+c)+2] = image[3*(r*w+c1)+2];
image[3*(r*w+c1)] = tmp[0];
image[3*(r*w+c1)+1] = tmp[1];
image[3*(r*w+c1)+2] = tmp[2];
}
}
}
// vertical flip
if (vflip) {
if (nchannel != 3) {
throw mjCError(
this, "currently only 3-channel textures support vertical flip");
}
for (int r=0; r < h/2; r++) {
for (int c=0; c < w; c++) {
int r1 = h-1-r;
unsigned char tmp[3] = {
image[3*(r*w+c)],
image[3*(r*w+c)+1],
image[3*(r*w+c)+2]
};
image[3*(r*w+c)] = image[3*(r1*w+c)];
image[3*(r*w+c)+1] = image[3*(r1*w+c)+1];
image[3*(r*w+c)+2] = image[3*(r1*w+c)+2];
image[3*(r1*w+c)] = tmp[0];
image[3*(r1*w+c)+1] = tmp[1];
image[3*(r1*w+c)+2] = tmp[2];
}
}
}
FlipIfNeeded(image, w, h);
}
// load 2D
void mjCTexture::Load2D(std::string filename, const mjVFS* vfs) {
// load PNG or custom
unsigned int w, h;
bool is_srgb;
std::vector<unsigned char> image;
LoadFlip(filename, vfs, image, w, h, is_srgb);
LoadFlip(filename, vfs, data_, w, h, is_srgb);
// assign size
width = w;
@@ -5077,24 +5072,8 @@ void mjCTexture::Load2D(std::string filename, const mjVFS* vfs) {
if (colorspace == mjCOLORSPACE_AUTO) {
colorspace = is_srgb ? mjCOLORSPACE_SRGB : mjCOLORSPACE_LINEAR;
}
// allocate and copy data
std::int64_t size = static_cast<std::int64_t>(width)*height;
if (size >= std::numeric_limits<int>::max() / nchannel || size <= 0) {
throw mjCError(this, "Texture too large");
}
try {
data_.assign(nchannel*size, std::byte(0));
} catch (const std::bad_alloc& e) {
throw mjCError(this, "Could not allocate memory for texture '%s' (id %d)",
(const char*)file_.c_str(), id);
}
memcpy(data_.data(), image.data(), nchannel*size);
image.clear();
}
// load cube or skybox from single file (repeated or grid)
void mjCTexture::LoadCubeSingle(std::string filename, const mjVFS* vfs) {
// check gridsize
@@ -5105,7 +5084,7 @@ void mjCTexture::LoadCubeSingle(std::string filename, const mjVFS* vfs) {
// load PNG or custom
unsigned int w, h;
bool is_srgb;
std::vector<unsigned char> image;
std::vector<std::byte> image;
LoadFlip(filename, vfs, image, w, h, is_srgb);
if (colorspace == mjCOLORSPACE_AUTO) {
@@ -5226,7 +5205,7 @@ void mjCTexture::LoadCubeSeparate(const mjVFS* vfs) {
// load PNG or custom
unsigned int w, h;
bool is_srgb;
std::vector<unsigned char> image;
std::vector<std::byte> image;
LoadFlip(filename.Str(), vfs, image, w, h, is_srgb);
// assume all faces have the same colorspace
@@ -5308,6 +5287,9 @@ void mjCTexture::Compile(const mjVFS* vfs) {
throw mjCError(this, "Texture buffer has incorrect size, given %d expected %d", nullptr,
data_.size(), nchannel * width * height);
}
// Flip if specified.
FlipIfNeeded(data_, width, height);
return;
}
src/user/user_objects.h
+8 -9
View File
@@ -1450,18 +1450,17 @@ class mjCTexture : public mjCTexture_, private mjsTexture {
void LoadCubeSingle(std::string filename, const mjVFS* vfs); // load cube from single file
void LoadCubeSeparate(const mjVFS* vfs); // load cube from separate files
void LoadFlip(std::string filename, const mjVFS* vfs, // load and flip
std::vector<unsigned char>& image,
unsigned int& w, unsigned int& h, bool& is_srgb);
void FlipIfNeeded(std::vector<std::byte>& image, unsigned int w, unsigned int h);
void LoadPNG(mjResource* resource,
std::vector<unsigned char>& image,
void LoadFlip(std::string filename, const mjVFS* vfs, // load and flip
std::vector<std::byte>& image, unsigned int& w, unsigned int& h,
bool& is_srgb);
void LoadPNG(mjResource* resource, std::vector<std::byte>& image,
unsigned int& w, unsigned int& h, bool& is_srgb);
void LoadKTX(mjResource* resource,
std::vector<unsigned char>& image,
void LoadKTX(mjResource* resource, std::vector<std::byte>& image,
unsigned int& w, unsigned int& h, bool& is_srgb);
void LoadCustom(mjResource* resource,
std::vector<unsigned char>& image,
void LoadCustom(mjResource* resource, std::vector<std::byte>& image,
unsigned int& w, unsigned int& h, bool& is_srgb);
bool clear_data_; // if true, data_ is empty and should be filled by Compile
test/user/user_api_test.cc
+53
View File
@@ -794,6 +794,59 @@ TEST_F(PluginTest, TextureFromBuffer) {
mj_deleteSpec(spec);
}
TEST_F(MujocoTest, TestTextureFlip) {
mjSpec* spec = mj_makeSpec();
EXPECT_THAT(spec, NotNull());
std::vector<std::byte> texture_data = {std::byte{1}, std::byte{2},
std::byte{3}, std::byte{4}};
mjsTexture* texture = mjs_addTexture(spec);
mjs_setName(texture->element, "hflipped");
texture->type = mjTEXTURE_2D;
texture->width = 2;
texture->height = 2;
texture->nchannel = 1;
texture->hflip = true;
mjs_setBuffer(texture->data, texture_data.data(), texture_data.size());
texture_data = {std::byte{11}, std::byte{12}, std::byte{13}, std::byte{21},
std::byte{22}, std::byte{23}, std::byte{31}, std::byte{32},
std::byte{33}, std::byte{41}, std::byte{42}, std::byte{43}};
texture = mjs_addTexture(spec);
mjs_setName(texture->element, "vflipped");
texture->type = mjTEXTURE_2D;
texture->width = 2;
texture->height = 2;
texture->nchannel = 3;
texture->vflip = true;
mjs_setBuffer(texture->data, texture_data.data(), texture_data.size());
mjModel* model = mj_compile(spec, 0);
EXPECT_THAT(model, NotNull()) << mjs_getError(spec);
EXPECT_THAT(model->tex_data[0], 2);
EXPECT_THAT(model->tex_data[1], 1);
EXPECT_THAT(model->tex_data[2], 4);
EXPECT_THAT(model->tex_data[3], 3);
EXPECT_THAT(model->tex_data[4], 31);
EXPECT_THAT(model->tex_data[5], 32);
EXPECT_THAT(model->tex_data[6], 33);
EXPECT_THAT(model->tex_data[7], 41);
EXPECT_THAT(model->tex_data[8], 42);
EXPECT_THAT(model->tex_data[9], 43);
EXPECT_THAT(model->tex_data[10], 11);
EXPECT_THAT(model->tex_data[11], 12);
EXPECT_THAT(model->tex_data[12], 13);
EXPECT_THAT(model->tex_data[13], 21);
EXPECT_THAT(model->tex_data[14], 22);
EXPECT_THAT(model->tex_data[15], 23);
mj_deleteModel(model);
mj_deleteSpec(spec);
}
// -------------------------------- test attach --------------------------------
static constexpr char xml_child[] = R"(