[MISC] Track visual overwrite by surface color as mesh metadata. (#2325)

Author: duburcqa

genesis/engine/mesh.py

@@ -62,9 +62,12 @@ def __init__(
         self._metadata: dict[str, Any] = metadata or {}
         self._color = np.array([1.0, 1.0, 1.0, 1.0], dtype=gs.np_float)
 
-        # By default, all meshes are considered zup, unless the "FileMorph.file_meshes_are_zup" option was set to False.
+        # By default, all meshes are considered zup, unless the "FileMorph.file_meshes_are_zup" option was set to False
         self._metadata.setdefault("imported_as_zup", True)
 
+        # By default, all meshes are considered having their original visual
+        self._metadata.setdefault("is_visual_overwritten", False)
+
         if not is_mesh_zup:
             if self._metadata["imported_as_zup"]:
                 self._mesh.apply_transform(mu.Y_UP_TRANSFORM.T)
@@ -245,6 +248,7 @@ def from_trimesh(
             # Visual may not have uv, e.g. ColorVisuals
             uvs = None
 
+        metadata = metadata or {}
         must_update_surface = True
         roughness_factor = None
         color_image = None
@@ -287,14 +291,14 @@ def from_trimesh(
             else:
                 # TODO: support vertex/face colors in luisa
                 color_factor = tuple(np.array(visual.main_color, dtype=np.float32) / 255.0)
+        elif surface.color is not None:
+            color_factor = surface.color
+            metadata["is_visual_overwritten"] = True
         elif (isinstance(visual, trimesh.visual.color.ColorVisuals) and visual.defined) or (
             isinstance(visual, trimesh.visual.color.VertexColor) and visual.vertex_colors.size > 0
         ):
-            # Color is already vertex-based. No need to create a new texture to keep the original one, unless a color
-            # overwrite has been specified as surface-level.
-            must_update_surface = surface.color is not None
-        elif surface.color is not None:
-            color_factor = surface.color
+            # Color is already vertex-based. It is not only necessary to create a new visual.
+            must_update_surface = False
         else:
             # use white color as default
             color_factor = (1.0, 1.0, 1.0, 1.0)
@@ -336,14 +340,20 @@ def from_attrs(
         if surface is None:
             surface = gs.surfaces.Default()
 
+        metadata = metadata or {}
+        metadata["is_visual_overwritten"] = metadata.get("is_visual_overwritten", False) or (surface.color is not None)
+        visual = mu.surface_uvs_to_trimesh_visual(surface, uvs, len(verts))
+
+        tmesh = trimesh.Trimesh(
+            vertices=verts,
+            faces=faces,
+            vertex_normals=normals,
+            visual=visual,
+            process=False,
+        )
+
         return cls(
-            mesh=trimesh.Trimesh(
-                vertices=verts,
-                faces=faces,
-                vertex_normals=normals,
-                visual=mu.surface_uvs_to_trimesh_visual(surface, uvs, len(verts)),
-                process=False,
-            ),
+            mesh=tmesh,
             surface=surface,
             uvs=uvs,
             scale=scale,
@@ -402,6 +412,7 @@ def update_trimesh_visual(self):
         Update the trimesh obj's visual attributes using its surface and uvs.
         """
         self._mesh.visual = mu.surface_uvs_to_trimesh_visual(self.surface, self.uvs, len(self.verts))
+        self._metadata["is_visual_overwritten"] = True
 
     def apply_transform(self, T):
         """

genesis/utils/urdf.py

@@ -117,6 +117,7 @@ def parse_urdf(morph, surface):
                     scale *= geometry.scale
 
                 # Overwrite surface color by original color specified in URDF file only if necessary
+                is_urdf_material = False
                 if geom_is_col:
                     geom_surface = gs.surfaces.Collision()
                 elif (
@@ -125,6 +126,7 @@ def parse_urdf(morph, surface):
                     and geom_prop.material.color is not None
                     and (morph.prioritize_urdf_material or surface.color is None)
                 ):
+                    is_urdf_material = True
                     geom_surface = surface.copy()
                     geom_surface.color = geom_prop.material.color
                 else:
@@ -145,6 +147,9 @@ def parse_urdf(morph, surface):
                         is_mesh_zup=morph.file_meshes_are_zup,
                         metadata={"mesh_path": mesh_path},
                     )
+                    if is_urdf_material:
+                        # Material color defined in URDF are not considered as visual overwrite
+                        mesh.metadata["is_visual_overwritten"] = False
 
                     g_info = {"mesh" if geom_is_col else "vmesh": mesh}
                     link_g_infos_.append(g_info)

tests/test_rigid_physics.py

@@ -2711,11 +2711,12 @@ def test_urdf_capsule(tmp_path, show_viewer, tol):
 @pytest.mark.required
 @pytest.mark.required
 @pytest.mark.parametrize("overwrite", [False, True])
-def test_urdf_color_overwrite(overwrite):
-    scene = gs.Scene()
+def test_urdf_color_overwrite(overwrite, show_viewer):
+    scene = gs.Scene(show_viewer=show_viewer)
     box = scene.add_entity(
         gs.morphs.URDF(
             file="genesis/assets/urdf/blue_box/model.urdf",
+            convexify=False,
         ),
         surface=gs.surfaces.Default(
             color=(1.0, 0.0, 0.0, 1.0) if overwrite else None,
@@ -2724,26 +2725,49 @@ def test_urdf_color_overwrite(overwrite):
     axis = scene.add_entity(
         gs.morphs.Mesh(
             file="meshes/axis.obj",
+            convexify=False,
         ),
         surface=gs.surfaces.Default(
             color=(1.0, 0.0, 0.0, 1.0) if overwrite else None,
         ),
     )
+    asset_path = get_hf_dataset(pattern="work_table.glb")
+    table = scene.add_entity(
+        gs.morphs.Mesh(
+            file=f"{asset_path}/work_table.glb",
+            convexify=False,
+        ),
+        surface=gs.surfaces.Default(
+            color=(1.0, 0.0, 0.0, 1.0) if overwrite else None,
+        ),
+    )
+    if show_viewer:
+        scene.build()
     for vgeom in box.vgeoms:
+        assert vgeom.vmesh.metadata["is_visual_overwritten"] == overwrite
         visual = vgeom.vmesh.trimesh.visual
         assert visual.defined
         color = np.unique(visual.vertex_colors, axis=0)
         assert_array_equal(color, (255, 0, 0, 255) if overwrite else (0, 0, 255, 255))
     for vgeom in axis.vgeoms:
+        assert vgeom.vmesh.metadata["is_visual_overwritten"] == overwrite
         visual = vgeom.vmesh.trimesh.visual
         assert visual.defined
         color = np.unique(visual.vertex_colors, axis=0)
         if overwrite:
             assert_array_equal(color, (255, 0, 0, 255))
         else:
             assert_array_equal(color, [[0, 0, 178, 255], [0, 178, 0, 255], [178, 0, 0, 255], [255, 255, 255, 255]])
+    for vgeom in table.vgeoms:
+        assert vgeom.vmesh.metadata["is_visual_overwritten"] == overwrite
+        visual = vgeom.vmesh.trimesh.visual
+        assert visual.defined
+        if overwrite:
+            color = np.unique(visual.vertex_colors, axis=0)
+            assert_array_equal(color, (255, 0, 0, 255))
     for entity in scene.entities:
         for geom in entity.geoms:
+            assert geom.mesh.metadata["is_visual_overwritten"]
             visual = geom.mesh.trimesh.visual
             assert visual.defined
             color = np.unique(visual.vertex_colors, axis=0)