[FEATURE] Add USD Materials Baking. (#1300)

Co-authored-by: Alexis Duburcq <alexis.duburcq@gmail.com>

Author: ACMLCZH

.github/workflows/production.yml

@@ -25,6 +25,8 @@ jobs:
       GENESIS_IMAGE_VER: "1_1"
       TIMEOUT_MINUTES: 180
       MADRONA_DISABLE_CUDA_HEAP_SIZE: "1"
+      OMNI_KIT_ACCEPT_EULA: "yes"
+      OMNI_KIT_ALLOW_ROOT: "1"
 
     steps:
       - name: Checkout code
@@ -48,6 +50,9 @@ jobs:
             NODELIST="--nodelist=$IDLE_NODES"
           fi
 
+          # TODO: USD baking does not currently support Python 3.11 since
+          # NVIDIA does not currently release `omniverse-kit==107.3` on PyPI.
+          # See: https://github.com/Genesis-Embodied-AI/Genesis/pull/1300
           srun \
             --container-image="/mnt/data/images/genesis-v${GENESIS_IMAGE_VER}.sqsh" \
             --container-mounts=\
@@ -60,7 +65,8 @@ jobs:
             --partition=hpc-mid ${NODELIST} --nodes=1 --time="${TIMEOUT_MINUTES}" \
             --job-name=${SLURM_JOB_NAME} \
             bash -c "
-              pip install -e '.[dev,render]' && \
+              pip install --extra-index-url https://pypi.nvidia.com/ omniverse-kit && \
+              pip install -e '.[dev,render,usd]' && \
               pytest -v --forked ./tests
             "
 

genesis/assets/tests/chopper.glb

@@ -13,7 +13,6 @@
 from genesis.options.surfaces import Surface
 import genesis.utils.mesh as mu
 import genesis.utils.gltf as gltf_utils
-import genesis.utils.usda as usda_utils
 import genesis.utils.particle as pu
 from genesis.repr_base import RBC
 
@@ -229,6 +228,7 @@ def from_trimesh(
         """
         if surface is None:
             surface = gs.surfaces.Default()
+            surface.update_texture()
         else:
             surface = surface.copy()
         mesh = mesh.copy(include_cache=True)
@@ -341,22 +341,20 @@ def from_morph_surface(cls, morph, surface=None):
         If the morph is a Mesh morph (morphs.Mesh), it could contain multiple submeshes, so we return a list.
         """
         if isinstance(morph, gs.options.morphs.Mesh):
-            if morph.file.endswith(("obj", "ply", "stl")):
+            if morph.is_format(gs.options.morphs.MESH_FORMATS):
                 meshes = mu.parse_mesh_trimesh(morph.file, morph.group_by_material, morph.scale, surface)
-
-            elif morph.file.endswith(("glb", "gltf")):
+            elif morph.is_format(gs.options.morphs.GLTF_FORMATS):
                 if morph.parse_glb_with_trimesh:
                     meshes = mu.parse_mesh_trimesh(morph.file, morph.group_by_material, morph.scale, surface)
                 else:
                     meshes = gltf_utils.parse_mesh_glb(morph.file, morph.group_by_material, morph.scale, surface)
+            elif morph.is_format(gs.options.morphs.USD_FORMATS):
+                import genesis.utils.usda as usda_utils
 
-            elif morph.file.endswith(("usd", "usda", "usdc", "usdz")):
                 meshes = usda_utils.parse_mesh_usd(morph.file, morph.group_by_material, morph.scale, surface)
-
             elif isinstance(morph, gs.options.morphs.MeshSet):
                 assert all(isinstance(mesh, trimesh.Trimesh) for mesh in morph.files)
                 meshes = [mu.trimesh_to_mesh(mesh, morph.scale, surface) for mesh in morph.files]
-
             else:
                 gs.raise_exception(
                     f"File type not supported (yet). Submit a feature request if you need this: {morph.file}."

genesis/assets/tests/combined_srt.glb

@@ -1,3 +1,10 @@
+"""
+We define all types of morphologies here: shape primitives, meshes, URDF, MJCF, and soft robot description files.
+
+These are independent of backend solver type and are shared by different solvers, e.g. a mesh can be either loaded as a
+rigid object / MPM object / FEM object.
+"""
+
 import os
 from typing import Any, List, Optional, Sequence, Tuple, Union
 
@@ -10,11 +17,12 @@
 from .misc import CoacdOptions
 from .options import Options
 
-"""
-We define all types of morphologies here: shape primitives, meshes, URDF, MJCF, and soft robot description files.
-These are independent of backend solver type and are shared by different solvers.
-E.g. a mesh can be either loaded as a rigid object / MPM object / FEM object.
-"""
+
+URDF_FORMAT = ".urdf"
+MJCF_FORMAT = ".xml"
+MESH_FORMATS = (".obj", ".ply", ".stl")
+GLTF_FORMATS = (".glb", ".gltf")
+USD_FORMATS = (".usd", ".usda", ".usdc", ".usdz")
 
 
 class TetGenMixin(Options):
@@ -556,6 +564,9 @@ def __init__(self, **data):
     def _repr_type(self):
         return f"<gs.morphs.{self.__class__.__name__}(file='{self.file}')>"
 
+    def is_format(self, format):
+        return self.file.lower().endswith(format)
+
 
 class Mesh(FileMorph, TetGenMixin):
     """
@@ -768,8 +779,8 @@ class MJCF(FileMorph):
 
     def __init__(self, **data):
         super().__init__(**data)
-        if not self.file.endswith(".xml"):
-            gs.raise_exception(f"Expected `.xml` extension for MJCF file: {self.file}")
+        if not self.is_format(MJCF_FORMAT):
+            gs.raise_exception(f"Expected `{MJCF_FORMAT}` extension for MJCF file: {self.file}")
 
         # What you want to do with scaling is kinda "zoom" the world from the perspective of the entity, i.e. scale the
         # geometric properties of an entity wrt its root pose. In the general case, ie for a 3D vector scale, (x, y, z)
@@ -878,8 +889,8 @@ class URDF(FileMorph):
 
     def __init__(self, **data):
         super().__init__(**data)
-        if isinstance(self.file, str) and not self.file.endswith(".urdf"):
-            gs.raise_exception(f"Expected `.urdf` extension for URDF file: {self.file}")
+        if isinstance(self.file, str) and not self.is_format(URDF_FORMAT):
+            gs.raise_exception(f"Expected `{URDF_FORMAT}` extension for URDF file: {self.file}")
 
         # Anisotropic scaling is ill-defined for poly-articulated robots. See related MJCF about this for details.
         if isinstance(self.scale, np.ndarray) and self.scale.std() > gs.EPS:
@@ -997,10 +1008,10 @@ def __init__(self, **data):
         # Make sure that Propellers links are preserved
         self.links_to_keep = tuple(set([*self.links_to_keep, *self.propellers_link_name]))
 
-        if isinstance(self.file, str) and not self.file.endswith(".urdf"):
-            gs.raise_exception(f"Drone only supports `.urdf` extension: {self.file}")
+        if isinstance(self.file, str) and not self.is_format(URDF_FORMAT):
+            gs.raise_exception(f"Drone only supports `{URDF_FORMAT}` extension: {self.file}")
 
-        if self.model not in ["CF2X", "CF2P", "RACE"]:
+        if self.model not in ("CF2X", "CF2P", "RACE"):
             gs.raise_exception(f"Unsupported `model`: {self.model}.")
 
 

genesis/assets/tests/combined_transform.glb

@@ -97,6 +97,9 @@ def __init__(self, **data):
             else:
                 gs.logger.warning("`env_euler` is ignored when `env_quat` is specified.")
 
+        if self.env_surface is not None:
+            self.env_surface.update_texture()
+
 
 class BatchRenderer(RendererOptions):
     """

genesis/engine/mesh.py

@@ -7,6 +7,7 @@
 from PIL import Image
 
 import genesis as gs
+import genesis.utils.mesh as mu
 
 from .options import Options
 
@@ -32,6 +33,9 @@ def check_simplify(self):
     def apply_cutoff(self, cutoff):
         raise NotImplementedError
 
+    def is_black(self):
+        raise NotImplementedError
+
 
 class ColorTexture(Texture):
     """
@@ -67,6 +71,9 @@ def apply_cutoff(self, cutoff):
             return
         self.color = tuple(1.0 if c >= cutoff else 0.0 for c in self.color)
 
+    def is_black(self):
+        return all(c < gs.EPS for c in self.color)
+
 
 class ImageTexture(Texture):
     """
@@ -115,25 +122,7 @@ def __init__(self, **data):
             if self.image_path.endswith((".hdr", ".exr")):
                 self.encoding = "linear"  # .exr or .hdr images should be encoded with 'linear'
                 if self.image_path.endswith((".exr")):
-                    exr_file = OpenEXR.InputFile(self.image_path)
-                    exr_header = exr_file.header()
-
-                    if exr_header["compression"].v > Imath.Compression.PIZ_COMPRESSION:
-                        new_image_path = f"{self.image_path[:-4]}_ZIP.exr"
-                        gs.logger.warning(
-                            f"EXR image {self.image_path}'s compression type {exr_header['compression']} is not supported. "
-                            f"Converting to compression type ZIP_COMPRESSION and saving to {new_image_path}."
-                        )
-                        self.image_path = new_image_path
-
-                        if not os.path.exists(new_image_path):
-                            channel_data = {channel: exr_file.channel(channel) for channel in exr_header["channels"]}
-                            exr_header["compression"] = Imath.Compression(Imath.Compression.ZIP_COMPRESSION)
-                            new_exr_file = OpenEXR.OutputFile(new_image_path, exr_header)
-                            new_exr_file.writePixels(channel_data)
-                            new_exr_file.close()
-
-                    exr_file.close()
+                    self.image_path = mu.check_exr_compression(self.image_path)
             else:
                 self.image_array = np.array(Image.open(self.image_path))
                 self.image_path = None
@@ -158,16 +147,8 @@ def __init__(self, **data):
                     )
             self.image_array = arr
 
-        # calculate channel
-        if self.image_array is None:
-            if isinstance(self.resolution, (tuple, list)):
-                H, W = self.resolution
-            else:
-                H = W = self.resolution
-
-            # Default to 3-channel RGB
-            white = np.array([255, 255, 255], dtype=np.uint8)
-            self.image_array = np.full((H, W, 3), white, dtype=np.uint8)
+        # just calculate channel
+        if self.image_array is None:  # Using 'image_path'
             self._mean_color = np.array([1.0, 1.0, 1.0], dtype=np.float16)
             self._channel = 3
         else:
@@ -216,3 +197,6 @@ def apply_cutoff(self, cutoff):
         if cutoff is None or self.image_array is None:  # Cutoff does not apply on image file.
             return
         self.image_array = np.where(self.image_array >= 255.0 * cutoff, 255, 0).astype(np.uint8)
+
+    def is_black(self):
+        return all(c < gs.EPS for c in self.image_color) or np.max(self.image_array) == 0

genesis/options/morphs.py

@@ -219,8 +219,9 @@ def parse_glb_material(glb, material_index, surface):
         if material.emissiveFactor is not None:
             emissive_factor = np.array(material.emissiveFactor, dtype=np.float32)
 
-        if emissive_factor is not None and np.any(emissive_factor > 0.0):  # Make sure to check emissive
-            emissive_texture = mu.create_texture(emissive_image, emissive_factor, "srgb")
+        emissive_texture = mu.create_texture(emissive_image, emissive_factor, "srgb")
+        if emissive_texture.is_black():  # Make sure to check emissive
+            emissive_texture = None
 
     # TODO: Parse them!
     for extension_name, extension_material in material.extensions.items():
@@ -317,7 +318,7 @@ def parse_mesh_glb(path, group_by_material, scale, surface):
                         primitive.material, parse_glb_material(glb, primitive.material, surface)
                     )
             else:
-                material, uv_used, material_name = None, 0, ""
+                material, uv_used, material_name = surface.copy(), 0, ""
 
             uvs = None
             if "KHR_draco_mesh_compression" in primitive.extensions:

genesis/options/renderers.py

@@ -1,13 +1,16 @@
 import hashlib
+import json
 import math
 import os
 import pickle as pkl
 from functools import lru_cache
-
+from pathlib import Path
 
 import numpy as np
 import trimesh
 from PIL import Image
+import OpenEXR
+import Imath
 
 import coacd
 import igl
@@ -20,10 +23,12 @@
 from .misc import (
     get_assets_dir,
     get_cvx_cache_dir,
+    get_exr_cache_dir,
     get_gsd_cache_dir,
     get_ptc_cache_dir,
     get_remesh_cache_dir,
     get_src_dir,
+    get_usd_cache_dir,
     get_tet_cache_dir,
 )
 
@@ -130,6 +135,26 @@ def get_remesh_path(verts, faces, edge_len_abs, edge_len_ratio, fix):
     return os.path.join(get_remesh_cache_dir(), f"{hashkey}.rm")
 
 
+def get_exr_path(file_path):
+    hashkey = get_file_hashkey(file_path)
+    return os.path.join(get_exr_cache_dir(), f"{hashkey}.exr")
+
+
+def get_usd_zip_path(file_path):
+    hashkey = get_file_hashkey(file_path)
+    return os.path.join(get_usd_cache_dir(), "zip", hashkey)
+
+
+def get_usd_bake_path(file_path):
+    hashkey = get_file_hashkey(file_path)
+    return os.path.join(get_usd_cache_dir(), "bake", hashkey)
+
+
+def get_file_hashkey(file):
+    file_obj = Path(file)
+    return get_hashkey(file_obj.resolve().as_posix().encode(), str(file_obj.stat().st_size).encode())
+
+
 def get_hashkey(*args):
     hasher = hashlib.sha256()
     for arg in args:
@@ -987,3 +1012,29 @@ def visualize_tet(tet, pv_data, show_surface=True, plot_cell_qual=False):
             plotter.add_mesh(pv_data, "r", "wireframe")
             plotter.add_legend([[" Input Mesh ", "r"], [" Tessellated Mesh ", "black"]])
             plotter.show()
+
+
+def check_exr_compression(exr_path):
+    exr_file = OpenEXR.InputFile(exr_path)
+    exr_header = exr_file.header()
+    if exr_header["compression"].v > Imath.Compression.PIZ_COMPRESSION:
+        new_exr_path = get_exr_path(exr_path)
+        if os.path.exists(new_exr_path):
+            gs.logger.info(f"Assets of fixed compression detected and used: {new_exr_path}.")
+        else:
+            gs.logger.warning(
+                f"EXR image {exr_path}'s compression type {exr_header['compression']} is not supported. "
+                f"Converting to compression type ZIP_COMPRESSION and saving to {new_exr_path}."
+            )
+
+            channel_data = {channel: exr_file.channel(channel) for channel in exr_header["channels"]}
+            exr_header["compression"] = Imath.Compression(Imath.Compression.ZIP_COMPRESSION)
+
+            os.makedirs(os.path.dirname(new_exr_path), exist_ok=True)
+            new_exr_file = OpenEXR.OutputFile(new_exr_path, exr_header)
+            new_exr_file.writePixels(channel_data)
+            new_exr_file.close()
+
+        exr_path = new_exr_path
+
+    exr_file.close()