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)
-----------------------------------
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
^^^^^^^^^^^^^^^
+3 -2
View File
@@ -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
+3 -2
View File
@@ -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
+2 -1
View File
@@ -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 ) \
+10 -3
View File
@@ -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',
+6 -5
View File
@@ -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 ) \
+3 -2
View File
@@ -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
)
)
+3 -3
View File
@@ -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);
}
}
+2 -1
View File
@@ -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;
+3 -3
View File
@@ -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());
+2 -1
View File
@@ -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;