Updates to material override

Author: cgokmen

omnigibson/examples/scenes/scene_selector.py

@@ -44,7 +44,15 @@ def main(random_selection=False, headless=False, short_exec=False):
         }
         load_mode = choose_from_options(options=load_options, name="load mode", random_selection=random_selection)
         if load_mode == "Quick":
-            cfg["scene"]["load_object_categories"] = ["floors", "walls", "ceilings"]
+            cfg["scene"]["load_object_categories"] = [
+                "floors",
+                "walls",
+                "ceilings",
+                "lawn",
+                "driveway",
+                "roof",
+                "rail_fence",
+            ]
 
     # Load the environment
     env = og.Environment(configs=cfg)

omnigibson/prims/entity_prim.py

@@ -15,15 +15,12 @@
 from omnigibson.prims.rigid_kinematic_prim import RigidKinematicPrim
 from omnigibson.prims.xform_prim import XFormPrim
 from omnigibson.utils.constants import JointAxis, JointType, PrimType
-from omnigibson.utils.ui_utils import create_module_logger
+from omnigibson.utils.render_utils import force_pbr_material_for_link
 from omnigibson.utils.usd_utils import PoseAPI, absolute_prim_path_to_scene_relative
 
 # Create settings for this module
 m = create_module_macros(module_path=__file__)
 
-# Create module logger
-log = create_module_logger(module_name=__name__)
-
 # Default sleep threshold for all objects -- see https://docs.omniverse.nvidia.com/extensions/latest/ext_physics/simulation-control/physics-settings.html?highlight=sleep#sleeping
 # Mass-normalized kinetic energy threshold below which an actor may go to sleep
 m.DEFAULT_SLEEP_THRESHOLD = 0.00005
@@ -262,25 +259,7 @@ def update_links(self):
 
             # Apply the V-Ray to PBR material change if request by the macro
             if gm.USE_PBR_MATERIALS:
-                # Find the material prim that has the link's name in it
-                link_pbr_material_pattern = f"__{link_name}_pbr"
-                looks_prim = self._prim.GetChild("Looks")
-                if not looks_prim:
-                    log.debug("Could not find Looks prim for", self.prim_path)
-                else:
-                    for mtl_prim in looks_prim.GetChildren():
-                        mtl_prim_name = mtl_prim.GetName()
-                        if link_pbr_material_pattern in mtl_prim_name:
-                            # Bind that material and stop
-                            lazy.omni.kit.commands.execute(
-                                "BindMaterialCommand",
-                                prim_path=prim.GetPrimPath().__str__(),
-                                material_path=mtl_prim.GetPrimPath().__str__(),
-                                strength=None,
-                            )
-                            break
-                    else:
-                        log.warning("Could not find PBR material for link", link_name, "of", self.prim_path)
+                force_pbr_material_for_link(self._prim, link_name)
 
             # Create and load the link
             self._links[link_name] = link_cls(

omnigibson/utils/render_utils.py

@@ -3,8 +3,12 @@
 """
 
 import omnigibson.lazy as lazy
-from omnigibson.prims import EntityPrim, RigidPrim, VisualGeomPrim
 from omnigibson.utils.physx_utils import bind_material
+from omnigibson.utils.ui_utils import create_module_logger
+
+
+# Create module logger
+log = create_module_logger(module_name=__name__)
 
 
 def make_glass(prim):
@@ -15,6 +19,9 @@ def make_glass(prim):
     Args:
         prim (EntityPrim or RigidPrim or VisualGeomPrim): Desired prim to convert into glass
     """
+    # Do this here to avoid circular imports
+    from omnigibson.prims import EntityPrim, RigidPrim, VisualGeomPrim
+
     # Generate the set of visual meshes we'll convert into glass
     if isinstance(prim, EntityPrim):
         # Grab all visual meshes from all links
@@ -72,3 +79,50 @@ def create_pbr_material(prim_path):
 
     # Return generated material
     return lazy.isaacsim.core.utils.prims.get_prim_at_path(material_path)
+
+
+def force_pbr_material_for_link(entity_prim, link_name):
+    if "meta__" in link_name:
+        # Don't override the meta link material
+        return
+
+    entity_prim_path = entity_prim.GetPrimPath().__str__()
+    looks_prim = entity_prim.GetChild("Looks")
+    link_prim = entity_prim.GetChild(link_name)
+    assert link_prim, f"Could not find link {link_name} for {entity_prim_path}"
+
+    link_visuals_prim = link_prim.GetChild("visuals")
+    if not looks_prim:
+        log.debug(f"Could not find Looks prim for {entity_prim_path}")
+        return
+
+    if not link_visuals_prim:
+        log.warning(f"Could not find visuals prim for link {link_name} of {entity_prim_path}")
+        return
+
+    binding_api = (
+        lazy.pxr.UsdShade.MaterialBindingAPI(link_visuals_prim)
+        if link_visuals_prim.HasAPI(lazy.pxr.UsdShade.MaterialBindingAPI)
+        else lazy.pxr.UsdShade.MaterialBindingAPI.Apply(link_visuals_prim)
+    )
+
+    material_path = binding_api.GetDirectBinding().GetMaterialPath().pathString
+    if "OmniGlass" in material_path:
+        # Don't override the glass material
+        return
+
+    # Find the material prim that has the link's name in it
+    link_pbr_material_pattern = f"__{link_name}_pbr"
+    for mtl_prim in looks_prim.GetChildren():
+        mtl_prim_name = mtl_prim.GetName()
+        if link_pbr_material_pattern in mtl_prim_name:
+            # Bind that material and stop
+            lazy.omni.kit.commands.execute(
+                "BindMaterialCommand",
+                prim_path=link_visuals_prim.GetPrimPath().__str__(),
+                material_path=mtl_prim.GetPrimPath().__str__(),
+                strength=None,
+            )
+            return
+    else:
+        log.warning(f"Could not find PBR material for link {link_name} of {entity_prim_path}")