Add mjmodel support for textures with different number of channels.

PiperOrigin-RevId: 653974959
Change-Id: Ia280e0bdb21027646249f104f21bedee792ba7e5
This commit is contained in:
Tom Erez
2024-07-19 05:21:16 -07:00
committed by Copybara-Service
parent 70ac76bb4b
commit e92af73cbf
11 changed files with 43 additions and 24 deletions
+6 -1
View File
@@ -5,10 +5,15 @@ Changelog
Upcoming version (not yet released) Upcoming version (not yet released)
----------------------------------- -----------------------------------
General
^^^^^^^
1. Renamed ``mjModel.tex_rbg`` to ``mjModel.tex_data``.
MJX 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 Python bindings
^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^
+3 -2
View File
@@ -1211,8 +1211,9 @@ struct mjModel_ {
int* tex_type; // texture type (mjtTexture) (ntex x 1) int* tex_type; // texture type (mjtTexture) (ntex x 1)
int* tex_height; // number of rows in texture image (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_width; // number of columns in texture image (ntex x 1)
int* tex_adr; // address in rgb (ntex x 1) int* tex_nchannel; // number of channels in texture image (ntex x 1)
mjtByte* tex_rgb; // rgb (alpha = 1) (ntexdata 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) int* tex_pathadr; // address of texture asset path; -1: none (ntex x 1)
// materials // materials
+3 -2
View File
@@ -924,8 +924,9 @@ struct mjModel_ {
int* tex_type; // texture type (mjtTexture) (ntex x 1) int* tex_type; // texture type (mjtTexture) (ntex x 1)
int* tex_height; // number of rows in texture image (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_width; // number of columns in texture image (ntex x 1)
int* tex_adr; // address in rgb (ntex x 1) int* tex_nchannel; // number of channels in texture image (ntex x 1)
mjtByte* tex_rgb; // rgb (alpha = 1) (ntexdata 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) int* tex_pathadr; // address of texture asset path; -1: none (ntex x 1)
// materials // materials
+2 -1
View File
@@ -410,8 +410,9 @@
X ( int, tex_type, ntex, 1 ) \ X ( int, tex_type, ntex, 1 ) \
X ( int, tex_height, ntex, 1 ) \ X ( int, tex_height, ntex, 1 ) \
X ( int, tex_width, ntex, 1 ) \ X ( int, tex_width, ntex, 1 ) \
X ( int, tex_nchannel, ntex, 1 ) \
X ( int, tex_adr, 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, tex_pathadr, ntex, 1 ) \
XMJV( int, mat_texid, nmat, mjNTEXMAT ) \ XMJV( int, mat_texid, nmat, mjNTEXMAT ) \
XMJV( mjtByte, mat_texuniform, nmat, 1 ) \ XMJV( mjtByte, mat_texuniform, nmat, 1 ) \
+10 -3
View File
@@ -2892,19 +2892,26 @@ STRUCTS: Mapping[str, StructDecl] = dict([
), ),
doc='number of columns in texture image (ntex x 1)', 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( StructFieldDecl(
name='tex_adr', name='tex_adr',
type=PointerType( type=PointerType(
inner_type=ValueType(name='int'), inner_type=ValueType(name='int'),
), ),
doc='address in rgb (ntex x 1)', doc='start address in tex_data (ntex x 1)',
), ),
StructFieldDecl( StructFieldDecl(
name='tex_rgb', name='tex_data',
type=PointerType( type=PointerType(
inner_type=ValueType(name='mjtByte'), inner_type=ValueType(name='mjtByte'),
), ),
doc='rgb (alpha = 1) (ntexdata x 1)', doc='pixel values (ntexdata x 1)',
), ),
StructFieldDecl( StructFieldDecl(
name='tex_pathadr', name='tex_pathadr',
+6 -5
View File
@@ -262,11 +262,12 @@
X( float, tendon, _rgba, ntendon, 4 ) X( float, tendon, _rgba, ntendon, 4 )
#define MJMODEL_TEXTURE \ #define MJMODEL_TEXTURE \
X( int, tex_, type, ntex, 1 ) \ X( int, tex_, type, ntex, 1 ) \
X( int, tex_, height, ntex, 1 ) \ X( int, tex_, height, ntex, 1 ) \
X( int, tex_, width, ntex, 1 ) \ X( int, tex_, width, ntex, 1 ) \
X( int, tex_, adr, ntex, 1 ) \ X( int, tex_, nchannel, ntex, 1 ) \
X( mjtByte, tex_, rgb, ntexdata, 1 ) X( int, tex_, adr, ntex, 1 ) \
X( mjtByte, tex_, data, ntexdata, 1 )
#define MJMODEL_TUPLE \ #define MJMODEL_TUPLE \
X( int, tuple_, adr, ntuple, 1 ) \ X( int, tuple_, adr, ntuple, 1 ) \
+3 -2
View File
@@ -212,9 +212,10 @@ class USDExporter:
for texture_id in tqdm.tqdm(range(self.model.ntex)): for texture_id in tqdm.tqdm(range(self.model.ntex)):
texture_height = self.model.tex_height[texture_id] texture_height = self.model.tex_height[texture_id]
texture_width = self.model.tex_width[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( 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 texture_height, texture_width, 3
) )
) )
+3 -3
View File
@@ -1382,7 +1382,7 @@ void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid) {
// assign data // assign data
glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, m->tex_width[texid], m->tex_height[texid], 0, 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 // generate mipmaps
glGenerateMipmap(GL_TEXTURE_2D); 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]) { if (m->tex_width[texid] == m->tex_height[texid]) {
for (int i=0; i < 6; i++) { for (int i=0; i < 6; i++) {
glTexImage2D(GL_TEXTURE_CUBE_MAP_POSITIVE_X+i, 0, GL_RGB, w, w, 0, 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 { else {
for (int i=0; i < 6; i++) { for (int i=0; i < 6; i++) {
glTexImage2D(GL_TEXTURE_CUBE_MAP_POSITIVE_X+i, 0, GL_RGB, w, w, 0, 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);
} }
} }
+2 -1
View File
@@ -2505,9 +2505,10 @@ void mjCModel::CopyObjects(mjModel* m) {
m->tex_height[i] = ptex->height; m->tex_height[i] = ptex->height;
m->tex_width[i] = ptex->width; m->tex_width[i] = ptex->width;
m->tex_adr[i] = data_adr; m->tex_adr[i] = data_adr;
m->tex_nchannel[i] = 3;
// copy rgb data // 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 // advance counter
data_adr += 3*ptex->width*ptex->height; data_adr += 3*ptex->width*ptex->height;
+3 -3
View File
@@ -288,15 +288,15 @@ TEST_F(PluginTest, RecompileComparePngCache) {
// load model once // load model once
mjModel* m = LoadModelFromString(xml, error.data(), error.size(), vfs.get()); 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 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); mj_deleteModel(m);
// update tex.png, load again // update tex.png, load again
mj_deleteFileVFS(vfs.get(), "tex.png"); mj_deleteFileVFS(vfs.get(), "tex.png");
mj_addBufferVFS(vfs.get(), "tex.png", tex2, sizeof(tex2)); mj_addBufferVFS(vfs.get(), "tex.png", tex2, sizeof(tex2));
m = LoadModelFromString(xml, error.data(), error.size(), vfs.get()); m = LoadModelFromString(xml, error.data(), error.size(), vfs.get());
EXPECT_NE(m->tex_rgb[0], byte); EXPECT_NE(m->tex_data[0], byte);
EXPECT_EQ(m->tex_rgb[15], byte); // first pixel is now last pixel EXPECT_EQ(m->tex_data[15], byte); // first pixel is now last pixel
mj_deleteModel(m); mj_deleteModel(m);
mj_deleteVFS(vfs.get()); mj_deleteVFS(vfs.get());
+2 -1
View File
@@ -5451,8 +5451,9 @@ public unsafe struct mjModel_ {
public int* tex_type; public int* tex_type;
public int* tex_height; public int* tex_height;
public int* tex_width; public int* tex_width;
public int* tex_nchannel;
public int* tex_adr; public int* tex_adr;
public byte* tex_rgb; public byte* tex_data;
public int* tex_pathadr; public int* tex_pathadr;
public int* mat_texid; public int* mat_texid;
public byte* mat_texuniform; public byte* mat_texuniform;