Add support for ORM maps in mujoco file format plugin

The ORM texture is a packed-channel map, with ambient occlusion in Red, roughness in Green, and metallic in Blue channel.

PiperOrigin-RevId: 772367369
Change-Id: I88a4dee2477daa1e3c0b00197bfbe2b276ba9937
This commit is contained in:
Google DeepMind
2025-06-17 01:44:17 -07:00
committed by Copybara-Service
parent e2e2932fe6
commit ea949f1cee
4 changed files with 104 additions and 23 deletions
@@ -92,7 +92,12 @@ TF_DEFINE_PRIVATE_TOKENS(kTokens,
((inputsWrapT, "inputs:wrapT"))
((inputsDiffuseColor, "inputs:diffuseColor"))
((outputsRgb, "outputs:rgb"))
((outputsR, "outputs:r"))
((outputsG, "outputs:g"))
((outputsB, "outputs:b"))
((inputsMetallic, "inputs:metallic"))
((inputsOcclusion, "inputs:occlusion"))
((inputsRoughness, "inputs:roughness"))
(repeat)
((sourceMesh, pxr::UsdGeomTokens->Mesh))
((inputsNormal, "inputs:normal"))
@@ -655,10 +660,10 @@ class ModelWriter {
return uvmap_st_output_attr;
}
pxr::SdfPath AddTextureShader(const pxr::SdfPath &material_path,
const char *texture_file,
const pxr::TfToken &name,
const pxr::SdfPath &uvmap_st_output_attr) {
std::vector<pxr::SdfPath> AddTextureShader(
const pxr::SdfPath &material_path, const char *texture_file,
const pxr::TfToken &name, const pxr::SdfPath &uvmap_st_output_attr,
const std::vector<pxr::TfToken> &output_channels) {
pxr::SdfPath texture_shader_path =
CreatePrimSpec(data_, material_path, name, pxr::UsdShadeTokens->Shader);
pxr::SdfPath texture_info_id_attr = CreateAttributeSpec(
@@ -688,11 +693,19 @@ class ModelWriter {
pxr::SdfValueTypeNames->Token);
SetAttributeDefault(data_, texture_wrap_t_attr, kTokens->repeat);
pxr::SdfPath texture_rgb_output_attr =
CreateAttributeSpec(data_, texture_shader_path, kTokens->outputsRgb,
pxr::SdfValueTypeNames->Float3);
return texture_rgb_output_attr;
std::vector<pxr::SdfPath> texture_output_attrs;
for (const auto &output_channel : output_channels) {
pxr::SdfValueTypeName value_type;
if (output_channel == kTokens->outputsRgb) {
value_type = pxr::SdfValueTypeNames->Float3;
} else {
// Assume the other specified channels are outputR, outputG, outputB.
value_type = pxr::SdfValueTypeNames->Float;
}
texture_output_attrs.push_back(CreateAttributeSpec(
data_, texture_shader_path, output_channel, value_type));
}
return texture_output_attrs;
}
void WriteMaterial(mjsMaterial *material, const pxr::SdfPath &parent_path) {
@@ -722,9 +735,44 @@ class ModelWriter {
const pxr::SdfPath &uvmap_st_output_attr =
AddUVTextureShader(material_path, pxr::TfToken("uvmap"));
// Find the normal texture if specified.
const mjStringVec &textures = *(material->textures);
// Set the values of metallic and roughness. These can come from an ORM
// texture or as a value defined in mjsMaterial_. Occlusion is only present
// in the ORM texture.
pxr::SdfPath metallic_attr = CreateAttributeSpec(
data_, preview_surface_shader_path, kTokens->inputsMetallic,
pxr::SdfValueTypeNames->Float);
pxr::SdfPath roughness_attr = CreateAttributeSpec(
data_, preview_surface_shader_path, kTokens->inputsRoughness,
pxr::SdfValueTypeNames->Float);
// Find the ORM (occlusion, roughness, metallic) packed-channel texture if
// specified.
if (mjTEXROLE_ORM < textures.size()) {
std::string orm_texture_name = textures[mjTEXROLE_ORM];
mjsTexture *orm_texture = mjs_asTexture(
mjs_findElement(spec_, mjOBJ_TEXTURE, orm_texture_name.c_str()));
if (orm_texture) {
// Create the ORM shader and connect its output to the preview
// surface ORM attrs.
const std::vector<pxr::SdfPath> orm_output_attrs = AddTextureShader(
material_path, orm_texture->file->c_str(),
pxr::TfToken("orm_packed"), uvmap_st_output_attr,
{kTokens->outputsR, kTokens->outputsG, kTokens->outputsB});
if (orm_output_attrs.size() == 3) {
pxr::SdfPath occlusion_attr = CreateAttributeSpec(
data_, preview_surface_shader_path, kTokens->inputsOcclusion,
pxr::SdfValueTypeNames->Float);
AddAttributeConnection(data_, occlusion_attr, orm_output_attrs[0]);
AddAttributeConnection(data_, roughness_attr, orm_output_attrs[1]);
AddAttributeConnection(data_, metallic_attr, orm_output_attrs[2]);
}
} else {
SetAttributeDefault(data_, metallic_attr, material->metallic);
SetAttributeDefault(data_, roughness_attr, material->roughness);
}
}
// Find the normal texture if specified.
if (mjTEXROLE_NORMAL < textures.size()) {
std::string normal_texture_name = textures[mjTEXROLE_NORMAL];
mjsTexture *normal_texture = mjs_asTexture(
@@ -735,10 +783,14 @@ class ModelWriter {
pxr::SdfValueTypeNames->Normal3f);
// Create the normal map shader and connect its output to the preview
// surface normal attr.
pxr::SdfPath normal_map_output_attr =
const std::vector<pxr::SdfPath> normal_map_output_attrs =
AddTextureShader(material_path, normal_texture->file->c_str(),
pxr::TfToken("normal"), uvmap_st_output_attr);
AddAttributeConnection(data_, normal_attr, normal_map_output_attr);
pxr::TfToken("normal"), uvmap_st_output_attr,
{kTokens->outputsRgb});
if (normal_map_output_attrs.size() == 1) {
AddAttributeConnection(data_, normal_attr,
normal_map_output_attrs[0]);
}
}
}
@@ -753,11 +805,14 @@ class ModelWriter {
if (main_texture) {
// Create the texture shader and connect it to the diffuse color
// attribute.
pxr::SdfPath texture_diffuse_output_attr =
const std::vector<pxr::SdfPath> texture_diffuse_output_attrs =
AddTextureShader(material_path, main_texture->file->c_str(),
pxr::TfToken("diffuse"), uvmap_st_output_attr);
AddAttributeConnection(data_, diffuse_color_attr,
texture_diffuse_output_attr);
pxr::TfToken("diffuse"), uvmap_st_output_attr,
{kTokens->outputsRgb});
if (texture_diffuse_output_attrs.size() == 1) {
AddAttributeConnection(data_, diffuse_color_attr,
texture_diffuse_output_attrs[0]);
}
} else {
// If no texture is specified, use the rgba diffuse color.
SetAttributeDefault(data_, diffuse_color_attr,
@@ -765,11 +820,6 @@ class ModelWriter {
material->rgba[2]));
}
pxr::SdfPath metallic_attr = CreateAttributeSpec(
data_, preview_surface_shader_path, kTokens->inputsMetallic,
pxr::SdfValueTypeNames->Float);
SetAttributeDefault(data_, metallic_attr, material->metallic);
pxr::SdfPath material_surface_output_attr = CreateAttributeSpec(
data_, material_path, pxr::UsdShadeTokens->outputsSurface,
pxr::SdfValueTypeNames->Token);
@@ -161,6 +161,7 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestMaterials) {
EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_layered");
EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_layered/uvmap");
EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_layered/diffuse");
EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_layered/normal");
ExpectAttributeHasConnection(
stage,
@@ -170,6 +171,33 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestMaterials) {
ExpectAttributeEqual(
stage, "/mesh_test/Materials/material_layered/normal.inputs:file",
pxr::SdfAssetPath("textures/normal.png"));
EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_layered/orm_packed");
ExpectAttributeHasConnection(
stage,
"/mesh_test/Materials/material_layered/"
"PreviewSurface.inputs:occlusion",
"/mesh_test/Materials/material_layered/orm_packed.outputs:r");
ExpectAttributeHasConnection(
stage,
"/mesh_test/Materials/material_layered/"
"PreviewSurface.inputs:roughness",
"/mesh_test/Materials/material_layered/orm_packed.outputs:g");
ExpectAttributeHasConnection(
stage,
"/mesh_test/Materials/material_layered/"
"PreviewSurface.inputs:metallic",
"/mesh_test/Materials/material_layered/orm_packed.outputs:b");
ExpectAttributeEqual(
stage, "/mesh_test/Materials/material_layered/orm_packed.inputs:file",
pxr::SdfAssetPath("textures/orm.png"));
EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_metallic");
EXPECT_PRIM_VALID(stage,
"/mesh_test/Materials/material_metallic/PreviewSurface");
ExpectAttributeEqual(stage,
"/mesh_test/Materials/material_metallic/PreviewSurface.inputs:metallic",
0.6f);
}
TEST_F(MjcfSdfFileFormatPluginTest, TestGeomRgba) {
@@ -2,11 +2,14 @@
<asset>
<texture name="diffuse_tex" file="textures/cube.png" type="2d" />
<texture name="normal_tex" file="textures/normal.png" type="2d" />
<texture name="orm_tex" file="textures/orm.png" type="2d" />
<material name="material_red" rgba="0.8 0 0 1" />
<material name="material_texture" texture="diffuse_tex" />
<material name="material_metallic" rgba="0.8 0 0 1" metallic="0.6"/>
<material name="material_layered">
<layer texture="diffuse_tex" role="rgb"/>
<layer texture="normal_tex" role="normal"/>
<layer texture="orm_tex" role="orm"/>
</material>
2d texture="diffuse" normal="normal" />
<mesh name="tetrahedron" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
Binary file not shown.

After

Width:  |  Height:  |  Size: 480 KiB