diff --git a/python/dist/mujoco-2.3.8.tar.gz b/python/dist/mujoco-2.3.8.tar.gz index b73593af..8db0e8a7 100644 Binary files a/python/dist/mujoco-2.3.8.tar.gz and b/python/dist/mujoco-2.3.8.tar.gz differ diff --git a/python/mujoco/usd_component.py b/python/mujoco/usd_component.py index 732fb902..fb242668 100644 --- a/python/mujoco/usd_component.py +++ b/python/mujoco/usd_component.py @@ -1,7 +1,8 @@ +import os from enum import Enum import mujoco -from pxr import Usd, UsdGeom, UsdLux, Vt, Gf +from pxr import Usd, UsdGeom, UsdLux, UsdShade, Vt, Gf, Sdf from scipy.spatial.transform import Rotation as R def create_usd_geom_primitive(geom, stage): @@ -146,7 +147,10 @@ class USDMesh(USDGeom): stage, model, mesh_vertex_ranges, - mesh_face_ranges): + mesh_face_ranges, + mesh_texcoord_ranges, + mesh_facetexcoord_ranges, + texture_file): super().__init__(geom, stage) self.type = 7 USDMesh.mesh_count += 1 @@ -164,12 +168,44 @@ class USDMesh(USDGeom): self.faces = model.mesh_face[mesh_face_ranges[mesh_idx]:mesh_face_ranges[mesh_idx+1]] self.prim.GetFaceVertexIndicesAttr().Set(self.faces) + texid = geom.texid + texcoords = model.mesh_texcoord[mesh_texcoord_ranges[texid]:mesh_texcoord_ranges[texid+1]] + + facetexcoords = model.mesh_facetexcoord.flatten() + facetexcoords = facetexcoords[mesh_facetexcoord_ranges[mesh_idx]:mesh_facetexcoord_ranges[mesh_idx+1]] + self.texcoords = UsdGeom.PrimvarsAPI(self.prim).CreatePrimvar("st", + Sdf.ValueTypeNames.TexCoord2fArray, + UsdGeom.Tokens.faceVarying) + + self.texcoords.Set(texcoords) + self.texcoords.SetIndices(Vt.IntArray(facetexcoords.tolist())); + + mtl_path = Sdf.Path(f"/World/Looks/Material_{os.path.splitext(os.path.basename(texture_file))[0]}") + mtl = UsdShade.Material.Define(stage, mtl_path) + shader = UsdShade.Shader.Define(stage, mtl_path.AppendPath("Shader")) + shader.CreateIdAttr("UsdPreviewSurface") + shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set((1.0, 0.0, 0.0)) + shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5) + shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0) + + diffuse_tx = UsdShade.Shader.Define(stage,mtl_path.AppendPath("DiffuseColorTx")) + diffuse_tx.CreateIdAttr('UsdUVTexture') + + # TODO: don't hardcode the image file + diffuse_tx.CreateInput('file', Sdf.ValueTypeNames.Asset).Set(texture_file) + diffuse_tx.CreateOutput('rgb', Sdf.ValueTypeNames.Float3) + shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(diffuse_tx.ConnectableAPI(), 'rgb') + mtl.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface") + + self.prim.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI) + UsdShade.MaterialBindingAPI(self.prim).Bind(mtl) + def update_geom(self, new_geom): self.update_pos(new_geom.pos) self.update_rotation(new_geom.mat) # TODO: remove this, temporary - self.xform.AddScaleOp().Set(value=(10.0, 10.0, 10.0)) + # self.xform.AddScaleOp().Set(value=(10.0, 10.0, 10.0)) class USDLight(object): """ diff --git a/python/mujoco/usd_renderer.py b/python/mujoco/usd_renderer.py index 08b94fb5..405fbe14 100644 --- a/python/mujoco/usd_renderer.py +++ b/python/mujoco/usd_renderer.py @@ -1,10 +1,12 @@ - +import os import mujoco import mujoco.viewer as viewer from mujoco.usd_component import * from mujoco.usd_utilities import * from pxr import Usd, UsdGeom +from PIL import Image as im + class USDRenderer(object): """ Renderer class the creates USD representations for mujoco scenes @@ -20,6 +22,7 @@ class USDRenderer(object): self.loaded_scene_info = False self.stage = Usd.Stage.CreateNew('usd_stage.usda') + UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z) @property def usd(self): @@ -52,9 +55,27 @@ class USDRenderer(object): if self.model.nmesh > 0: mesh_vertex_ranges = get_mesh_ranges(self.model.nmesh, self.model.mesh_vertnum) mesh_face_ranges = get_mesh_ranges(self.model.nmesh, self.model.mesh_facenum) + mesh_texcoord_ranges = get_mesh_ranges(self.model.nmesh, self.model.mesh_texcoordnum) + mesh_facetexcoord_ranges = get_facetexcoord_ranges(self.model.nmesh, self.model.mesh_facenum) + + + # create and load the texture files + # iterate through all the textures and build list of tex_rgb ranges + data_adr = 0 + texture_files = [] + for texid in range(self.model.ntex): + height = self.model.tex_height[texid] + width = self.model.tex_width[texid] + pixels = 3*height*width + rgb = self.model.tex_rgb[data_adr:data_adr+pixels] + img = rgb.reshape(height, width, 3) + file_name = f'{texid}.png' + im.fromarray(img).save(file_name) + texture_file = os.path.abspath(file_name) + texture_files.append(texture_file) + data_adr += pixels current_mesh_idx = 0 - # initializes an array to store all the geoms in the scene # populates with "empty" USDGeom objects self.usd_geoms = [] @@ -67,7 +88,10 @@ class USDRenderer(object): self.stage, self.model, mesh_vertex_ranges, - mesh_face_ranges)) + mesh_face_ranges, + mesh_texcoord_ranges, + mesh_facetexcoord_ranges, + texture_files[geoms[i].texid])) current_mesh_idx += 1 else: self.usd_geoms.append(create_usd_geom_primitive(geoms[i], self.stage)) @@ -110,7 +134,7 @@ class USDRenderer(object): def start_viewer(self): if self.data: - viewer.launch(self.model, self.data) + viewer.launch(self.model) def render(self): # should render the usd file given a particular renderer that diff --git a/python/mujoco/usd_utilities.py b/python/mujoco/usd_utilities.py index 6989c0b5..19ee07f6 100644 --- a/python/mujoco/usd_utilities.py +++ b/python/mujoco/usd_utilities.py @@ -5,4 +5,12 @@ def get_mesh_ranges(nmesh, arr): for i in range(nmesh): running_sum += arr[i] mesh_ranges.append(running_sum) - return mesh_ranges \ No newline at end of file + return mesh_ranges + +def get_facetexcoord_ranges(nmesh, arr): + facetexcoords_ranges = [0] + running_sum = 0 + for i in range(nmesh): + running_sum += arr[i] * 3 + facetexcoords_ranges.append(running_sum) + return facetexcoords_ranges \ No newline at end of file