diff --git a/.github/workflows/production.yml b/.github/workflows/production.yml index b56d92a123..07fd8d84ea 100644 --- a/.github/workflows/production.yml +++ b/.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 " diff --git a/genesis/assets/tests/chopper.glb b/genesis/assets/tests/chopper.glb deleted file mode 100644 index dbbac8a817..0000000000 Binary files a/genesis/assets/tests/chopper.glb and /dev/null differ diff --git a/genesis/assets/tests/combined_srt.glb b/genesis/assets/tests/combined_srt.glb deleted file mode 100644 index 7934814094..0000000000 Binary files a/genesis/assets/tests/combined_srt.glb and /dev/null differ diff --git a/genesis/assets/tests/combined_transform.glb b/genesis/assets/tests/combined_transform.glb deleted file mode 100644 index bd5b1d0bb7..0000000000 Binary files a/genesis/assets/tests/combined_transform.glb and /dev/null differ diff --git a/genesis/engine/mesh.py b/genesis/engine/mesh.py index 95f5d55fde..42f5eb6515 100644 --- a/genesis/engine/mesh.py +++ b/genesis/engine/mesh.py @@ -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}." diff --git a/genesis/options/morphs.py b/genesis/options/morphs.py index 4ef69d9d94..d46413065d 100644 --- a/genesis/options/morphs.py +++ b/genesis/options/morphs.py @@ -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"" + 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}.") diff --git a/genesis/options/renderers.py b/genesis/options/renderers.py index 65c483bcfe..3051c24fdd 100644 --- a/genesis/options/renderers.py +++ b/genesis/options/renderers.py @@ -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): """ diff --git a/genesis/options/textures.py b/genesis/options/textures.py index e5dda5f0b2..a16dcd90c5 100644 --- a/genesis/options/textures.py +++ b/genesis/options/textures.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 diff --git a/genesis/utils/gltf.py b/genesis/utils/gltf.py index 5104e05c03..d4b9d9ffb4 100644 --- a/genesis/utils/gltf.py +++ b/genesis/utils/gltf.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: diff --git a/genesis/utils/mesh.py b/genesis/utils/mesh.py index 1214c959b1..331fe79d03 100644 --- a/genesis/utils/mesh.py +++ b/genesis/utils/mesh.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() diff --git a/genesis/utils/misc.py b/genesis/utils/misc.py index 4280ddf415..f031fb6fdb 100644 --- a/genesis/utils/misc.py +++ b/genesis/utils/misc.py @@ -270,6 +270,14 @@ def get_remesh_cache_dir(): return os.path.join(get_cache_dir(), "rm") +def get_exr_cache_dir(): + return os.path.join(get_cache_dir(), "exr") + + +def get_usd_cache_dir(): + return os.path.join(get_cache_dir(), "usd") + + def clean_cache_files(): folder = gs.utils.misc.get_cache_dir() try: diff --git a/genesis/utils/urdf.py b/genesis/utils/urdf.py index f87ed42fea..f073d5c379 100644 --- a/genesis/utils/urdf.py +++ b/genesis/utils/urdf.py @@ -1,5 +1,6 @@ import os from itertools import chain +from pathlib import Path import trimesh import numpy as np @@ -54,7 +55,7 @@ def _order_links(l_infos, j_infos, links_g_infos=None): def parse_urdf(morph, surface): - if isinstance(morph.file, str): + if isinstance(morph.file, (str, Path)): path = os.path.join(get_assets_dir(), morph.file) robot = urdfpy.URDF.load(path) else: diff --git a/genesis/utils/usda.py b/genesis/utils/usda.py index 0a3fa516fa..dfeb58cc04 100644 --- a/genesis/utils/usda.py +++ b/genesis/utils/usda.py @@ -1,11 +1,26 @@ -import genesis as gs -from . import mesh as mu +import io +import os +import sys +import shutil +import subprocess +import logging +from pathlib import Path -from pxr import Usd, UsdGeom, UsdShade -import trimesh import numpy as np +import trimesh from PIL import Image -import io + +import genesis as gs + +from . import mesh as mu + +try: + from pxr import Usd, UsdGeom, UsdShade, Sdf +except ImportError as e: + gs.raise_exception_from( + "Failed to import USD dependencies. Try installing Genesis with 'usd' optional dependencies.", e + ) + cs_encode = { "raw": "linear", @@ -30,14 +45,7 @@ def get_input_attribute_value(shader, input_name, input_type=None): return None, None -def get_texture_image(image_path, zipfiles): - if zipfiles is None: - return np.array(Image.open(image_path.resolvedPath)) - else: - return np.array(Image.open(io.BytesIO(zipfiles.GetFile(image_path.path)))) - - -def parse_preview_surface(shader, output_name, zipfiles): +def parse_preview_surface(shader, output_name): shader_id = shader.GetShaderId() if shader_id == "UsdPreviewSurface": uvname = None @@ -50,7 +58,7 @@ def parse_component(component_name, component_encode): component_uvname = None else: # texture shader component_image, component_overencode, component_uvname = parse_preview_surface( - component, component_output, zipfiles + component, component_output ) if component_overencode is not None: component_encode = component_overencode @@ -74,6 +82,8 @@ def parse_component(component_name, component_encode): # parse emissive emissive_texture, emissive_uvname = parse_component("emissiveColor", "srgb") + if emissive_texture is not None and emissive_texture.is_black(): + emissive_texture = None if emissive_uvname is not None and uvname is None: uvname = emissive_uvname @@ -96,9 +106,12 @@ def parse_component(component_name, component_encode): if normal_uvname is not None and uvname is None: uvname = normal_uvname - # parse io + # parse ior ior = get_input_attribute_value(shader, "ior", "value")[0] + if uvname is None: + uvname = "st" + return { "color_texture": color_texture, "opacity_texture": opacity_texture, @@ -112,7 +125,7 @@ def parse_component(component_name, component_encode): elif shader_id == "UsdUVTexture": texture = get_input_attribute_value(shader, "file", "value")[0] if texture is not None: - texture_image = get_texture_image(texture, zipfiles) + texture_image = np.asarray(Image.open(texture.resolvedPath)) if output_name == "r": texture_image = texture_image[:, :, 0] elif output_name == "g": @@ -131,7 +144,7 @@ def parse_component(component_name, component_encode): texture_encode = get_input_attribute_value(shader, "sourceColorSpace", "value")[0] or "sRGB" texture_encode = cs_encode[texture_encode] texture_uvs_shader, texture_uvs_output = get_input_attribute_value(shader, "st", "attribute") - texture_uvs_name = parse_preview_surface(texture_uvs_shader, texture_uvs_output, zipfiles) + texture_uvs_name = parse_preview_surface(texture_uvs_shader, texture_uvs_output) return texture_image, texture_encode, texture_uvs_name @@ -140,163 +153,30 @@ def parse_component(component_name, component_encode): return primvar_name -def parse_gltf_surface(shader, source_type, output_name, zipfiles): - shader_subid = shader.GetSourceAssetSubIdentifier(source_type) - if shader_subid == "gltf_material": - # Parse color - color_factor = get_input_attribute_value(shader, "base_color_factor", "value")[0] # Gf.Vec3f(1.0, 1.0, 1.0) - color_texture_shader, color_texture_output = get_input_attribute_value( - shader, "base_color_texture", "attribute" - ) - if color_texture_shader is not None: - color_image = parse_gltf_surface(color_texture_shader, source_type, color_texture_output, zipfiles) - else: - color_image = None - color_texture = mu.create_texture(color_image, color_factor, "srgb") - - # parse opacity - opacity_factor = get_input_attribute_value(shader, "base_alpha", "value")[0] - opacity_texture = mu.create_texture(None, opacity_factor, "linear") - alpha_cutoff = get_input_attribute_value(shader, "alpha_cutoff", "value")[0] - alpha_mode = get_input_attribute_value(shader, "alpha_mode", "value")[0] - alpha_cutoff = mu.adjust_alpha_cutoff(alpha_cutoff, alpha_mode) - opacity_texture.apply_cutoff(alpha_cutoff) - - # parse roughness and metaillic - metallic_factor = get_input_attribute_value(shader, "metallic_factor", "value")[0] - roughness_factor = get_input_attribute_value(shader, "roughness_factor", "value")[0] - combined_texture_shader, combined_texture_output = get_input_attribute_value( - shader, "metallic_roughness_texture", "attribute" - ) - if combined_texture_shader is not None: - combined_image = parse_gltf_surface(combined_texture_shader, source_type, combined_texture_output, zipfiles) - roughness_image = combined_image[:, :, 1] - metallic_image = combined_image[:, :, 2] - else: - roughness_image, metallic_image = None, None - metallic_texture = mu.create_texture(metallic_image, metallic_factor, "linear") - roughness_texture = mu.create_texture(roughness_image, roughness_factor, "linear") - - # parse emissive - emissive_strength = get_input_attribute_value(shader, "emissive_strength", "value")[0] - emissive_texture = mu.create_texture(None, emissive_strength, "srgb") if emissive_strength else None - - occlusion_texture_shader, occlusion_texture_output = get_input_attribute_value( - shader, "occlusion_texture", "attribute" - ) - if occlusion_texture_shader is not None: - occlusion_image = parse_gltf_surface( - occlusion_texture_shader, source_type, occlusion_texture_output, zipfiles - ) - - return { - "color_texture": color_texture, - "opacity_texture": opacity_texture, - "roughness_texture": roughness_texture, - "metallic_texture": metallic_texture, - "emissive_texture": emissive_texture, - }, "st" - - elif shader_subid == "gltf_texture_lookup": - texture = get_input_attribute_value(shader, "texture", "value")[0] - if texture is not None: - texture_image = get_texture_image(texture, zipfiles) - else: - texture_image = None - return texture_image - - else: - raise Exception(f"Fail to parse gltf Shader {shader_subid}.") - - -def parse_omni_surface(shader, source_type, output_name, zipfiles): - - def parse_component(component_name, component_encode, adjust=None): - component_usetex = get_input_attribute_value(shader, f"Is{component_name}Tex", "value")[0] == 1 - if component_usetex: - component_tex_name = f"{component_name}_Tex" - component_tex = get_input_attribute_value(shader, component_tex_name, "value")[0] - if component_tex is not None: - component_image = get_texture_image(component_tex, zipfiles) - if adjust is not None: - component_image = (adjust(component_image / 255.0) * 255.0).astype(np.uint8) - component_cs = shader.GetInput(component_tex_name).GetAttr().GetColorSpace() - component_overencode = cs_encode[component_cs] - if component_overencode is not None: - component_encode = component_overencode - component_factor = None - else: - component_color_name = f"{component_name}_Color" - component_factor = get_input_attribute_value(shader, component_color_name, "value")[0] - if adjust is not None and component_factor is not None: - component_factor = tuple([adjust(c) for c in component_factor]) - component_image = None - - component_texture = mu.create_texture(component_image, component_factor, component_encode) - return component_texture - - color_texture = parse_component("BaseColor", "srgb") - opacity_texture = color_texture.check_dim(3) if color_texture else None - emissive_intensity = get_input_attribute_value(shader, "EmissiveIntensity", "value")[0] - emissive_texture = ( - parse_component("Emissive", "srgb", lambda x: x * emissive_intensity) if emissive_intensity else None - ) - if emissive_texture is not None: - emissive_texture.check_dim(3) - metallic_texture = parse_component("Metallic", "linear") - normal_texture = parse_component("Normal", "linear") - roughness_texture = parse_component("Gloss", "linear", lambda x: (2 / (x + 2)) ** (1.0 / 4.0)) - - return { - "color_texture": color_texture, - "opacity_texture": opacity_texture, - "roughness_texture": roughness_texture, - "metallic_texture": metallic_texture, - "emissive_texture": emissive_texture, - "normal_texture": normal_texture, - }, "st" - - -def parse_usd_material(material, surface, zipfiles): +def parse_usd_material(material, surface): surface_outputs = material.GetSurfaceOutputs() material_dict, uv_name = None, None - material_surface = None + material_surface = surface.copy() + + require_bake = False + material_candidates = [] for surface_output in surface_outputs: if not surface_output.HasConnectedSource(): continue surface_output_connectable, surface_output_name, _ = surface_output.GetConnectedSource() surface_shader = UsdShade.Shader(surface_output_connectable.GetPrim()) surface_shader_implement = surface_shader.GetImplementationSource() + surface_shader_id = surface_shader.GetShaderId() - if surface_shader_implement == "id": - shader_id = surface_shader.GetShaderId() - if shader_id == "UsdPreviewSurface": - material_dict, uv_name = parse_preview_surface(surface_shader, surface_output_name, zipfiles) - break - gs.logger.warning(f"Fail to parse Shader {surface_shader.GetPath()} with ID {shader_id}.") - continue + if surface_shader_implement == "id" and surface_shader_id == "UsdPreviewSurface": + material_dict, uv_name = parse_preview_surface(surface_shader, surface_output_name) + require_bake = False + break - elif surface_shader_implement == "sourceAsset": - source_types = surface_shader.GetSourceTypes() - for source_type in source_types: - source_asset = surface_shader.GetSourceAsset(source_type).resolvedPath - if "gltf/pbr" in source_asset: - material_dict, uv_name = parse_gltf_surface( - surface_shader, source_type, surface_output_name, zipfiles - ) - break - try: - material_dict, uv_name = parse_omni_surface( - surface_shader, source_type, surface_output_name, zipfiles - ) - except Exception as e: - gs.logger.warning( - f"Fail to parse Shader {surface_shader.GetPath()} of asset {source_asset} with message: {e}." - ) - continue + material_candidates.append((surface_shader.GetPath(), surface_shader_id, surface_shader_implement)) + require_bake = True if material_dict is not None: - material_surface = surface.copy() material_surface.update_texture( color_texture=material_dict.get("color_texture"), opacity_texture=material_dict.get("opacity_texture"), @@ -306,19 +186,161 @@ def parse_usd_material(material, surface, zipfiles): emissive_texture=material_dict.get("emissive_texture"), ior=material_dict.get("ior"), ) - return material_surface, uv_name + + if require_bake: + candidates_str = "\n".join( + f"\tShader at {shader_path} with implement {shader_impl} and ID {shader_id}." + for shader_path, shader_id, shader_impl in material_candidates + ) + gs.logger.debug(f"Material require baking:\n{candidates_str}") + return material_surface, uv_name, require_bake -def parse_mesh_usd(path, group_by_material, scale, surface): - zipfiles = Usd.ZipFile.Open(path) if path.endswith(".usdz") else None +def replace_asset_symlinks(stage): + asset_paths = set() + + for prim in stage.TraverseAll(): + for attr in prim.GetAttributes(): + value = attr.Get() + if isinstance(value, Sdf.AssetPath): + asset_paths.add(value.resolvedPath) + elif isinstance(value, list): + for v in value: + if isinstance(v, Sdf.AssetPath): + asset_paths.add(v.resolvedPath) + + for asset_path in map(Path, asset_paths): + if not asset_path.is_symlink(): + continue + + real_path = asset_path.resolve() + if asset_path.suffix.lower() == real_path.suffix.lower(): + continue + + asset_path.unlink() + if real_path.is_file(): + gs.logger.warning(f"Replacing symlink {asset_path} with real file {real_path}.") + shutil.copy2(real_path, asset_path) + + +def decompress_usdz(usdz_path): + usdz_folder = mu.get_usd_zip_path(usdz_path) + + # The first file in the package must be a native usd file. + # See https://openusd.org/docs/Usdz-File-Format-Specification.html + zip_files = Usd.ZipFile.Open(usdz_path) + zip_filelist = zip_files.GetFileNames() + root_file = zip_filelist[0] + if not root_file.lower().endswith(gs.options.morphs.USD_FORMATS[:-1]): + gs.raise_exception(f"Invalid usdz root file: {root_file}") + root_path = os.path.join(usdz_folder, root_file) + + if not os.path.exists(root_path): + for file_name in zip_filelist: + file_data = io.BytesIO(zip_files.GetFile(file_name)) + file_path = os.path.join(usdz_folder, file_name) + file_folder = os.path.dirname(file_path) + os.makedirs(file_folder, exist_ok=True) + with open(file_path, "wb") as out: + out.write(file_data.read()) + gs.logger.warning(f"USDZ file {usdz_path} decompressed to {root_path}.") + else: + gs.logger.info(f"Decompressed assets detected and used: {root_path}.") + return root_path + + +def parse_mesh_usd(path, group_by_material, scale, surface, bake_cache=True): + if path.lower().endswith(gs.options.morphs.USD_FORMATS[-1]): + path = decompress_usdz(path) + + # detect bake file caches + is_bake_cache_found = False + baked_folder = mu.get_usd_bake_path(path) + baked_path = os.path.join(baked_folder, os.path.basename(path)) + if bake_cache and os.path.exists(baked_path): + path = baked_path + is_bake_cache_found = True + gs.logger.info(f"Baked assets detected and used: {path}") + stage = Usd.Stage.Open(path) scale *= UsdGeom.GetStageMetersPerUnit(stage) yup = UsdGeom.GetStageUpAxis(stage) == "Y" xform_cache = UsdGeom.XformCache() mesh_infos = mu.MeshInfoGroup() - materials = dict() + materials = {} + baked_materials = {} + # parse materials + for prim in stage.Traverse(): + if prim.IsA(UsdShade.Material): + material_usd = UsdShade.Material(prim) + material_spec = prim.GetPrimStack()[-1] + material_id = material_spec.layer.identifier + material_spec.path.pathString + material_pack = materials.get(material_id, None) + + if material_pack is None: + material, uv_name, require_bake = parse_usd_material(material_usd, surface) + materials[material_id] = (material, uv_name) + if not is_bake_cache_found and require_bake: + baked_materials[material_id] = material_usd.GetPath() + + if baked_materials: + device = gs.device + if device.type == "cpu": + try: + device, *_ = gs.utils.get_device(gs.cuda) + except gs.GenesisException as e: + gs.raise_exception_from("USD baking requires CUDA GPU.", e) + + replace_asset_symlinks(stage) + os.makedirs(baked_folder, exist_ok=True) + + # Note that it is necessary to call 'bake_usd_material' via a subprocess to ensure proper isolation of + # omninerse kit, otherwise the global conversion registry of some Python bindings will be conflicting between + # each, ultimately leading to segfault... + commands = [ + "python", + os.path.join(os.path.dirname(os.path.abspath(__file__)), "usda_bake.py"), + "--input_file", + path, + "--output_dir", + baked_folder, + "--usd_material_paths", + *map(str, baked_materials.values()), + "--device", + str(device.index), + "--log_level", + logging.getLevelName(gs.logger.level).lower(), + ] + gs.logger.debug(f"Execute: {' '.join(commands)}") + + try: + result = subprocess.run( + commands, + capture_output=True, + check=True, + text=True, + ) + if result.stdout: + gs.logger.debug(result.stdout) + if result.stderr: + gs.logger.warning(result.stderr) + except (subprocess.CalledProcessError, OSError) as e: + gs.logger.warning(f"Baking process failed: {e} (Note that USD baking may only support Python 3.10 now.)") + + if os.path.exists(baked_path): + gs.logger.warning(f"USD materials baked to file {baked_path}") + stage = Usd.Stage.Open(baked_path) + for baked_material_id, baked_material_path in baked_materials.items(): + baked_material_usd = UsdShade.Material(stage.GetPrimAtPath(baked_material_path)) + baked_material, uv_name, require_bake = parse_usd_material(baked_material_usd, surface) + materials[baked_material_id] = (baked_material, uv_name) + + for baked_texture_obj in Path(baked_folder).glob("baked_textures*"): + shutil.rmtree(baked_texture_obj) + + # parse geometries for prim in stage.Traverse(): if prim.HasRelationship("material:binding"): if not prim.HasAPI(UsdShade.MaterialBindingAPI): @@ -355,26 +377,25 @@ def parse_mesh_usd(path, group_by_material, scale, surface): material_usd = prim_bindings.ComputeBoundMaterial()[0] if material_usd.GetPrim().IsValid(): material_spec = material_usd.GetPrim().GetPrimStack()[-1] - material_id = material_spec.layer.identifier + material_spec.path.pathString + material_file = material_spec.layer.identifier + material_file = path if material_file == baked_path else material_file + material_id = material_file + material_spec.path.pathString material, uv_name = materials.get(material_id, (None, "st")) - if material is None: - material, uv_name = materials.setdefault( - material_id, parse_usd_material(material_usd, surface, zipfiles) - ) else: - material, uv_name, material_id = None, "st", None + material, uv_name, material_id = surface.copy(), "st", None # parse uvs - uv_var = UsdGeom.PrimvarsAPI(prim).GetPrimvar(uv_name) uvs = None - if uv_var.IsDefined() and uv_var.HasValue(): - uvs = np.array(uv_var.ComputeFlattened(), dtype=np.float32) - if uvs.shape[0] != points.shape[0]: - if uvs.shape[0] == faces.shape[0]: - points_faces_varying = True - else: - gs.raise_exception(f"Size of uvs mismatch for mesh {mesh_id} in usd file {path}.") - uvs[:, 1] = 1.0 - uvs[:, 1] + if uv_name is not None: + uv_var = UsdGeom.PrimvarsAPI(prim).GetPrimvar(uv_name) + if uv_var.IsDefined() and uv_var.HasValue(): + uvs = np.array(uv_var.ComputeFlattened(), dtype=np.float32) + if uvs.shape[0] != points.shape[0]: + if uvs.shape[0] == faces.shape[0]: + points_faces_varying = True + else: + gs.raise_exception(f"Size of uvs mismatch for mesh {mesh_id} in usd file {path}.") + uvs[:, 1] = 1.0 - uvs[:, 1] # rearrange points and faces if points_faces_varying: @@ -402,13 +423,14 @@ def parse_mesh_usd(path, group_by_material, scale, surface): vertices=points, faces=triangles, vertex_normals=normals, - visual=trimesh.visual.TextureVisuals(uv=uvs), + visual=trimesh.visual.TextureVisuals(uv=uvs) if uvs is not None else None, process=True, ) points = processed_mesh.vertices triangles = processed_mesh.faces normals = processed_mesh.vertex_normals - uvs = processed_mesh.visual.uv + if uvs is not None: + uvs = processed_mesh.visual.uv # apply tranform points, normals = mu.apply_transform(matrix, points, normals) @@ -417,7 +439,13 @@ def parse_mesh_usd(path, group_by_material, scale, surface): mesh_info, first_created = mesh_infos.get(group_idx) if first_created: mesh_info.set_property( - surface=material, metadata={"path": path, "name": material_id if group_by_material else mesh_id} + surface=material, + metadata={ + "path": path, # unbaked file or cache + "name": material_id if group_by_material else mesh_id, + "require_bake": material_id in baked_materials, + "bake_success": material_id in baked_materials and material is not None, + }, ) mesh_info.append(points, triangles, normals, uvs) diff --git a/genesis/utils/usda_bake.py b/genesis/utils/usda_bake.py new file mode 100644 index 0000000000..f2f9213dc7 --- /dev/null +++ b/genesis/utils/usda_bake.py @@ -0,0 +1,115 @@ +import argparse +import asyncio +import logging +import os +import time + + +def omni_bootstrap(device=0, log_level="warning"): + import omni.kit_app + + app = omni.kit_app.KitApp() + kit_dir = os.path.dirname(os.path.abspath(os.path.realpath(omni.kit_app.__file__))) + kit_path = os.path.join(kit_dir, "apps", "omni.app.empty.kit") + app_args = [ + kit_path, + "--/app/window/hideUi=True", + f"--/app/tokens/exe-path={kit_dir}", + "--/app/enableStdoutOutput=False", # Disable print outs (print_and_log) on extension startup information + "--/app/runLoops/present/rateLimitFrequency=60", + "--/app/vulkan=True", + "--/app/asyncRendering=False", + "--/app/python/interceptSysStdOutput=False", + "--/app/python/logSysStdOutput=False", + "--/app/settings/fabricDefaultStageFrameHistoryCount=3", + f"--/omni/log/level={log_level}", + "--/log/file=", # Empty string means no log file + f"--/log/level={log_level}", + f"--/renderer/activeGpu={device}", + "--/renderer/enabled=rtx", + "--/renderer/active=rtx", + "--/renderer/multiGpu/enabled=False", # Avoids unnecessary GPU context initialization + "--no-window", + "--portable", + "--enable", + "omni.usd", + "--enable", + "omni.kit.material.library", + "--enable", + "omni.kit.viewport.utility", + "--enable", + "omni.kit.viewport.rtx", + "--enable", + "omni.kit.usd.collect", + "--enable", + "omni.replicator.core", + "--enable", + "omni.mdl.distill_and_bake", + ] + app.startup(app_args) + app.update() # important + return app + + +def bake_usd_material(input_file, output_dir, usd_material_paths, device=0, log_level="error"): + logs = [] + + # bootstrap + start_time = time.time() + app = omni_bootstrap(device, log_level) + logs.append(f"\tBootstrap: {time.time() - start_time}, App status: {app.is_running()}.") + + import omni.usd + import omni.mdl.distill_and_bake + import omni.replicator.core + import omni.kit.usd.collect + + # open stage + start_time = time.time() + omni.usd.get_context().open_stage(input_file) + logs.append(f"\tOpen stage: {time.time() - start_time}s.") + + # create render product + start_time = time.time() + stage = omni.usd.get_context().get_stage() + render_prod_path = omni.replicator.core.create.render_product("/OmniverseKit_Persp", resolution=(600, 600)) + app.update() # important + logs.append(f"\tCreate render product: {time.time() - start_time}s, {render_prod_path}.") + + # distill the material + start_time = time.time() + + for usd_material_path in usd_material_paths: + material_prim = stage.GetPrimAtPath(usd_material_path) + distiller = omni.mdl.distill_and_bake.MdlDistillAndBake(material_prim, ouput_folder=output_dir) + distiller.distill() + logs.append(f"\tDistill: {time.time() - start_time}s, {material_prim}.") + + # export usd + start_time = time.time() + collector = omni.kit.usd.collect.Collector(stage.GetRootLayer().identifier, output_dir) + task = asyncio.ensure_future(collector.collect()) + while not task.done(): + app.update() # Otherwise it will be blocked by omni.kit.app.get_app().next_update_async() + success, baked_file = task.result() + + logs.append(f"\tExport: {time.time() - start_time}s, {baked_file}.") + print("Distill USD material:\n" + "\n".join(logs)) + + # close omniverse app + app.shutdown() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--input_file", type=str, required=True) + parser.add_argument("--output_dir", type=str, required=True) + parser.add_argument("--usd_material_paths", type=str, nargs="+", required=True) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--log_level", type=str, default="warning") + args = parser.parse_args() + + log_level = logging.getLevelName( + min(max(logging.getLevelName(args.log_level.upper()), logging.INFO), logging.ERROR) + ).lower() + bake_usd_material(args.input_file, args.output_dir, args.usd_material_paths, args.device, log_level) diff --git a/genesis/vis/raytracer.py b/genesis/vis/raytracer.py index 6ef6a89def..893bdcb88a 100644 --- a/genesis/vis/raytracer.py +++ b/genesis/vis/raytracer.py @@ -157,19 +157,11 @@ def __init__(self, options, vis_options): self.lights = [] for light in options.lights: light_intensity = light.get("intensity", 1.0) - self.lights.append( - SphereLight( - radius=light["radius"], - pos=light["pos"], - surface=gs.surfaces.Emission( - color=( - light["color"][0] * light_intensity, - light["color"][1] * light_intensity, - light["color"][2] * light_intensity, - ), - ), - ) + light_surface = gs.surfaces.Emission( + color=map(lambda x: x * light_intensity, light["color"]), ) + light_surface.update_texture() + self.lights.append(SphereLight(radius=light["radius"], pos=light["pos"], surface=light_surface)) LuisaRenderPy.init( context_path=LRP_PATH, diff --git a/pyproject.toml b/pyproject.toml index 89e10594a1..3653265305 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,6 +100,10 @@ render = [ "pybind11[global]", "open3d", ] +usd = [ + # Used for parsing `.usd` mesh files + "usd-core", +] [project.scripts] gs = "genesis._main:main" diff --git a/tests/test_mesh.py b/tests/test_mesh.py index 85f807c782..bb3d1bc8fd 100644 --- a/tests/test_mesh.py +++ b/tests/test_mesh.py @@ -1,5 +1,7 @@ import os +import sys from pathlib import Path + import numpy as np import pytest import trimesh @@ -11,6 +13,7 @@ from .utils import assert_allclose, assert_array_equal, get_hf_dataset + VERTICES_TOL = 1e-05 # Transformation loses a little precision in vertices NORMALS_TOL = 1e-02 # Conversion from .usd to .glb loses a little precision in normals @@ -144,10 +147,11 @@ def check_gs_textures(gs_texture1, gs_texture2, default_value, material_name, te @pytest.mark.required -@pytest.mark.parametrize("glb_file", ["tests/combined_srt.glb", "tests/combined_transform.glb"]) +@pytest.mark.parametrize("glb_file", ["glb/combined_srt.glb", "glb/combined_transform.glb"]) def test_glb_parse_geometry(glb_file): """Test glb mesh geometry parsing.""" - glb_file = os.path.join(mu.get_assets_dir(), glb_file) + asset_path = get_hf_dataset(pattern=glb_file) + glb_file = os.path.join(asset_path, glb_file) gs_meshes = gltf_utils.parse_mesh_glb( glb_file, group_by_material=False, @@ -171,10 +175,11 @@ def test_glb_parse_geometry(glb_file): @pytest.mark.required -@pytest.mark.parametrize("glb_file", ["tests/chopper.glb"]) +@pytest.mark.parametrize("glb_file", ["glb/chopper.glb"]) def test_glb_parse_material(glb_file): """Test glb mesh geometry parsing.""" - glb_file = os.path.join(mu.get_assets_dir(), glb_file) + asset_path = get_hf_dataset(pattern=glb_file) + glb_file = os.path.join(asset_path, glb_file) gs_meshes = gltf_utils.parse_mesh_glb( glb_file, group_by_material=True, @@ -294,27 +299,63 @@ def test_usd_parse(usd_filename): ) +@pytest.mark.skipif( + sys.version_info[:2] != (3, 10) or sys.platform not in ("linux", "win32"), + reason="omniverse-kit used by USD Baking cannot be correctly installed on this platform now.", +) +@pytest.mark.parametrize( + "usd_file", ["usd/WoodenCrate/WoodenCrate_D1_1002.usda", "usd/franka_mocap_teleop/table_scene.usd"] +) +def test_usd_bake(usd_file, show_viewer): + asset_path = get_hf_dataset(pattern=os.path.join(os.path.dirname(usd_file), "*"), local_dir_use_symlinks=False) + usd_file = os.path.join(asset_path, usd_file) + gs_usd_meshes = usda_utils.parse_mesh_usd( + usd_file, group_by_material=True, scale=1.0, surface=gs.surfaces.Default(), bake_cache=False + ) + for gs_usd_mesh in gs_usd_meshes: + require_bake = gs_usd_mesh.metadata["require_bake"] + bake_success = gs_usd_mesh.metadata["bake_success"] + assert not require_bake or (require_bake and bake_success) + + scene = gs.Scene( + show_viewer=show_viewer, + show_FPS=False, + ) + robot = scene.add_entity( + gs.morphs.Mesh( + file=usd_file, + ), + ) + scene.build() + + @pytest.mark.required def test_urdf_with_existing_glb(tmp_path, show_viewer): - assets = Path(gs.utils.get_assets_dir()) - glb_path = assets / "usd" / "sneaker_airforce.glb" + glb_file = "usd/sneaker_airforce.glb" + asset_path = get_hf_dataset(pattern=glb_file) + urdf_path = tmp_path / "model.urdf" urdf_path.write_text( f""" - + """ ) + scene = gs.Scene( show_viewer=show_viewer, show_FPS=False, ) + robot = scene.add_entity( + gs.morphs.URDF( + file=urdf_path, + ), + ) scene.build() - scene.step() @pytest.mark.required @@ -348,7 +389,6 @@ def test_urdf_with_float_texture_glb(tmp_path, show_viewer, n_channels, float_ty glb_path = tmp_path / f"tex_{n_channels}c.glb" urdf_path = tmp_path / f"tex_{n_channels}c.urdf" trimesh.Scene([mesh]).export(glb_path) - urdf_path.write_text( f""" @@ -359,6 +399,11 @@ def test_urdf_with_float_texture_glb(tmp_path, show_viewer, n_channels, float_ty """ ) + scene = gs.Scene(show_viewer=show_viewer, show_FPS=False) + robot = scene.add_entity( + gs.morphs.URDF( + file=urdf_path, + ), + ) scene.build() - scene.step() diff --git a/tests/utils.py b/tests/utils.py index afdd71ddda..bb13ef4b7d 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -24,14 +24,16 @@ from genesis.utils import mjcf as mju from genesis.utils.mesh import get_assets_dir from genesis.utils.misc import tensor_to_array +from genesis.options.morphs import URDF_FORMAT, MJCF_FORMAT, MESH_FORMATS, GLTF_FORMATS, USD_FORMATS REPOSITY_URL = "Genesis-Embodied-AI/Genesis" DEFAULT_BRANCH_NAME = "main" -MESH_EXTENSIONS = (".mtl", ".glb", ".obj", ".stl", ".usb", ".usdz", ".mdl") +MESH_EXTENSIONS = (".mtl", *MESH_FORMATS, *GLTF_FORMATS, *USD_FORMATS) IMAGE_EXTENSIONS = (".png", ".jpg") + # Get repository "root" path (actually test dir is good enough) TEST_DIR = os.path.dirname(__file__) @@ -169,7 +171,12 @@ def get_git_commit_info(ref="HEAD"): def get_hf_dataset( - pattern, repo_name: str = "assets", local_dir: str | None = None, num_retry: int = 4, retry_delay: float = 30.0 + pattern, + repo_name: str = "assets", + local_dir: str | None = None, + num_retry: int = 4, + retry_delay: float = 30.0, + local_dir_use_symlinks: bool = True, ): assert num_retry >= 1 @@ -183,6 +190,7 @@ def get_hf_dataset( allow_patterns=pattern, max_workers=1, local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, ) # Make sure that download was successful @@ -192,7 +200,7 @@ def get_hf_dataset( continue ext = path.suffix.lower() - if not ext in (".xml", ".urdf", *IMAGE_EXTENSIONS, *MESH_EXTENSIONS): + if not ext in (URDF_FORMAT, MJCF_FORMAT, *IMAGE_EXTENSIONS, *MESH_EXTENSIONS): continue has_files = True @@ -200,7 +208,7 @@ def get_hf_dataset( if path.stat().st_size == 0: raise HTTPError(f"File '{path}' is empty.") - if path.suffix.lower() in (".xml", ".urdf"): + if path.suffix.lower() in (URDF_FORMAT, MJCF_FORMAT): try: ET.parse(path) except ET.ParseError as e: