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

PiperOrigin-RevId: 813256285
Change-Id: I0d49885dc483ab8f093b2566ec4bef355575505e

Author: havess

src/user/user_objects.cc

@@ -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,93 +5055,25 @@ 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;
   height = h;
   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

@@ -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

@@ -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"(