From e92af73cbfe30cfd7355ac8e428f38de330a1abe Mon Sep 17 00:00:00 2001 From: Tom Erez Date: Fri, 19 Jul 2024 05:21:16 -0700 Subject: [PATCH] Add mjmodel support for textures with different number of channels. PiperOrigin-RevId: 653974959 Change-Id: Ia280e0bdb21027646249f104f21bedee792ba7e5 --- doc/changelog.rst | 7 ++++++- doc/includes/references.h | 5 +++-- include/mujoco/mjmodel.h | 5 +++-- include/mujoco/mjxmacro.h | 3 ++- introspect/structs.py | 13 ++++++++++--- python/mujoco/indexer_xmacro.h | 11 ++++++----- python/mujoco/usd/exporter.py | 5 +++-- src/render/render_context.c | 6 +++--- src/user/user_model.cc | 3 ++- test/user/user_api_test.cc | 6 +++--- unity/Runtime/Bindings/MjBindings.cs | 3 ++- 11 files changed, 43 insertions(+), 24 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 88918ff0..ab2bdbdc 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -5,10 +5,15 @@ Changelog Upcoming version (not yet released) ----------------------------------- +General +^^^^^^^ + +1. Renamed ``mjModel.tex_rbg`` to ``mjModel.tex_data``. + MJX ^^^ -1. Added more fields to ``mjx.Model`` and ``mjx.Data`` for further compatibility with the corresponding MuJoCo structs. +2. Added more fields to ``mjx.Model`` and ``mjx.Data`` for further compatibility with the corresponding MuJoCo structs. Python bindings ^^^^^^^^^^^^^^^ diff --git a/doc/includes/references.h b/doc/includes/references.h index 17735c3b..b96978f9 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -1211,8 +1211,9 @@ struct mjModel_ { int* tex_type; // texture type (mjtTexture) (ntex x 1) int* tex_height; // number of rows in texture image (ntex x 1) int* tex_width; // number of columns in texture image (ntex x 1) - int* tex_adr; // address in rgb (ntex x 1) - mjtByte* tex_rgb; // rgb (alpha = 1) (ntexdata x 1) + int* tex_nchannel; // number of channels in texture image (ntex x 1) + int* tex_adr; // start address in tex_data (ntex x 1) + mjtByte* tex_data; // pixel values (ntexdata x 1) int* tex_pathadr; // address of texture asset path; -1: none (ntex x 1) // materials diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 2bafdd13..2262932e 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -924,8 +924,9 @@ struct mjModel_ { int* tex_type; // texture type (mjtTexture) (ntex x 1) int* tex_height; // number of rows in texture image (ntex x 1) int* tex_width; // number of columns in texture image (ntex x 1) - int* tex_adr; // address in rgb (ntex x 1) - mjtByte* tex_rgb; // rgb (alpha = 1) (ntexdata x 1) + int* tex_nchannel; // number of channels in texture image (ntex x 1) + int* tex_adr; // start address in tex_data (ntex x 1) + mjtByte* tex_data; // pixel values (ntexdata x 1) int* tex_pathadr; // address of texture asset path; -1: none (ntex x 1) // materials diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index cb9dc58d..4d33d9e6 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -410,8 +410,9 @@ X ( int, tex_type, ntex, 1 ) \ X ( int, tex_height, ntex, 1 ) \ X ( int, tex_width, ntex, 1 ) \ + X ( int, tex_nchannel, ntex, 1 ) \ X ( int, tex_adr, ntex, 1 ) \ - X ( mjtByte, tex_rgb, ntexdata, 1 ) \ + X ( mjtByte, tex_data, ntexdata, 1 ) \ XMJV( int, tex_pathadr, ntex, 1 ) \ XMJV( int, mat_texid, nmat, mjNTEXMAT ) \ XMJV( mjtByte, mat_texuniform, nmat, 1 ) \ diff --git a/introspect/structs.py b/introspect/structs.py index 5bdf71cd..b2a4f893 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -2892,19 +2892,26 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='number of columns in texture image (ntex x 1)', ), + StructFieldDecl( + name='tex_nchannel', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of channels in texture image (ntex x 1)', + ), StructFieldDecl( name='tex_adr', type=PointerType( inner_type=ValueType(name='int'), ), - doc='address in rgb (ntex x 1)', + doc='start address in tex_data (ntex x 1)', ), StructFieldDecl( - name='tex_rgb', + name='tex_data', type=PointerType( inner_type=ValueType(name='mjtByte'), ), - doc='rgb (alpha = 1) (ntexdata x 1)', + doc='pixel values (ntexdata x 1)', ), StructFieldDecl( name='tex_pathadr', diff --git a/python/mujoco/indexer_xmacro.h b/python/mujoco/indexer_xmacro.h index 0d0b529f..d729384c 100644 --- a/python/mujoco/indexer_xmacro.h +++ b/python/mujoco/indexer_xmacro.h @@ -262,11 +262,12 @@ X( float, tendon, _rgba, ntendon, 4 ) #define MJMODEL_TEXTURE \ - X( int, tex_, type, ntex, 1 ) \ - X( int, tex_, height, ntex, 1 ) \ - X( int, tex_, width, ntex, 1 ) \ - X( int, tex_, adr, ntex, 1 ) \ - X( mjtByte, tex_, rgb, ntexdata, 1 ) + X( int, tex_, type, ntex, 1 ) \ + X( int, tex_, height, ntex, 1 ) \ + X( int, tex_, width, ntex, 1 ) \ + X( int, tex_, nchannel, ntex, 1 ) \ + X( int, tex_, adr, ntex, 1 ) \ + X( mjtByte, tex_, data, ntexdata, 1 ) #define MJMODEL_TUPLE \ X( int, tuple_, adr, ntuple, 1 ) \ diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 0d2702c3..e3c126eb 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -212,9 +212,10 @@ class USDExporter: for texture_id in tqdm.tqdm(range(self.model.ntex)): texture_height = self.model.tex_height[texture_id] texture_width = self.model.tex_width[texture_id] - pixels = 3 * texture_height * texture_width + texture_nchannel = self.model.tex_nchannel[texture_id] + pixels = texture_nchannel * texture_height * texture_width img = im.fromarray( - self.model.tex_rgb[data_adr : data_adr + pixels].reshape( + self.model.tex_data[data_adr : data_adr + pixels].reshape( texture_height, texture_width, 3 ) ) diff --git a/src/render/render_context.c b/src/render/render_context.c index a4f0a4a8..96f863de 100644 --- a/src/render/render_context.c +++ b/src/render/render_context.c @@ -1382,7 +1382,7 @@ void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid) { // assign data glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, m->tex_width[texid], m->tex_height[texid], 0, - GL_RGB, GL_UNSIGNED_BYTE, m->tex_rgb + m->tex_adr[texid]); + GL_RGB, GL_UNSIGNED_BYTE, m->tex_data + m->tex_adr[texid]); // generate mipmaps glGenerateMipmap(GL_TEXTURE_2D); @@ -1416,7 +1416,7 @@ void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid) { if (m->tex_width[texid] == m->tex_height[texid]) { for (int i=0; i < 6; i++) { glTexImage2D(GL_TEXTURE_CUBE_MAP_POSITIVE_X+i, 0, GL_RGB, w, w, 0, - GL_RGB, GL_UNSIGNED_BYTE, m->tex_rgb + m->tex_adr[texid]); + GL_RGB, GL_UNSIGNED_BYTE, m->tex_data + m->tex_adr[texid]); } } @@ -1424,7 +1424,7 @@ void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid) { else { for (int i=0; i < 6; i++) { glTexImage2D(GL_TEXTURE_CUBE_MAP_POSITIVE_X+i, 0, GL_RGB, w, w, 0, - GL_RGB, GL_UNSIGNED_BYTE, m->tex_rgb + m->tex_adr[texid] + i*3*w*w); + GL_RGB, GL_UNSIGNED_BYTE, m->tex_data + m->tex_adr[texid] + i*3*w*w); } } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index a096bf7c..6517414b 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -2505,9 +2505,10 @@ void mjCModel::CopyObjects(mjModel* m) { m->tex_height[i] = ptex->height; m->tex_width[i] = ptex->width; m->tex_adr[i] = data_adr; + m->tex_nchannel[i] = 3; // copy rgb data - memcpy(m->tex_rgb + data_adr, ptex->rgb.data(), 3*ptex->width*ptex->height); + memcpy(m->tex_data + data_adr, ptex->rgb.data(), 3*ptex->width*ptex->height); // advance counter data_adr += 3*ptex->width*ptex->height; diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 1e811104..093c07c1 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -288,15 +288,15 @@ TEST_F(PluginTest, RecompileComparePngCache) { // load model once mjModel* m = LoadModelFromString(xml, error.data(), error.size(), vfs.get()); EXPECT_EQ(m->ntexdata, 18); // w x h x rgb = 3 x 2 x 3 - mjtByte byte = m->tex_rgb[0]; + mjtByte byte = m->tex_data[0]; mj_deleteModel(m); // update tex.png, load again mj_deleteFileVFS(vfs.get(), "tex.png"); mj_addBufferVFS(vfs.get(), "tex.png", tex2, sizeof(tex2)); m = LoadModelFromString(xml, error.data(), error.size(), vfs.get()); - EXPECT_NE(m->tex_rgb[0], byte); - EXPECT_EQ(m->tex_rgb[15], byte); // first pixel is now last pixel + EXPECT_NE(m->tex_data[0], byte); + EXPECT_EQ(m->tex_data[15], byte); // first pixel is now last pixel mj_deleteModel(m); mj_deleteVFS(vfs.get()); diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 635be2dc..4dd43872 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -5451,8 +5451,9 @@ public unsafe struct mjModel_ { public int* tex_type; public int* tex_height; public int* tex_width; + public int* tex_nchannel; public int* tex_adr; - public byte* tex_rgb; + public byte* tex_data; public int* tex_pathadr; public int* mat_texid; public byte* mat_texuniform;