From 58dc62c4ecfdbe500a23c83e639a22e8f0f078ab Mon Sep 17 00:00:00 2001
From: Abhishek Joshi <abhishekjoshinj@gmail.com>
Date: Tue, 13 Aug 2024 11:09:02 -0700
Subject: [PATCH] Fixing USD pipeline after copybara changes

---
 python/mujoco/usd/objects.py | 10 ++++++++--
 python/mujoco/usd/shapes.py  | 10 +++++-----
 python/mujoco/usd/utils.py   |  1 +
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py
index 850e3645..f74eae6b 100644
--- a/python/mujoco/usd/objects.py
+++ b/python/mujoco/usd/objects.py
@@ -392,8 +392,14 @@ class USDPrimitiveMesh(USDObject):
         if self.geom.size[1] > 0:
           t_scale *= self.geom.size[1]
 
-      mesh_texcoord[:, 0] *= s_scale / (self.geom.size[0] * 2)
-      mesh_texcoord[:, 1] *= t_scale / (self.geom.size[1] * 2)
+      s_size, t_size = self.geom.size[:2]
+      if self.geom.type == mujoco.mjtGeom.mjGEOM_PLANE:
+        s_size = s_size if s_size > 0 else 1
+        t_size = t_size if t_size > 0 else 1
+
+      if self.model.mat_texuniform[self.geom.matid]:
+        mesh_texcoord[:, 0] *= s_scale / (s_size * 2)
+        mesh_texcoord[:, 1] *= t_scale / (t_size * 2)
 
     return mesh_texcoord, mesh_facetexcoord.flatten()
 
diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py
index ca29b276..ceaea43b 100644
--- a/python/mujoco/usd/shapes.py
+++ b/python/mujoco/usd/shapes.py
@@ -26,7 +26,7 @@ def get_triangle_uvs(
     texture_type: Optional[mujoco.mjtTexture]
 ):
   """Returns UV coordinates for a given mesh."""
-  if not texture_type:
+  if texture_type is None:
     return None
 
   triangle_uvs = []
@@ -110,7 +110,7 @@ class TriangleMesh:
       height: float,
       depth: float,
       texture_type: Optional[mujoco.mjtTexture]
-  ) -> TriangleMesh:
+  ) -> "TriangleMesh":
     """Creates a box."""
     vertices = np.array([[0.0, 0.0, 0.0],
                          [width, 0.0, 0.0],
@@ -144,7 +144,7 @@ class TriangleMesh:
       radius: float,
       texture_type: Optional[mujoco.mjtTexture],
       resolution: int
-  ) -> TriangleMesh:
+  ) -> "TriangleMesh":
     """Creates a sphere."""
     vertices = []
     triangles = []
@@ -178,7 +178,7 @@ class TriangleMesh:
       radius: float,
       texture_type: Optional[mujoco.mjtTexture],
       resolution: int,
-  ) -> TriangleMesh:
+  ) -> "TriangleMesh":
     """Creates a hemisphere."""
     vertices = []
     triangles = []
@@ -218,7 +218,7 @@ class TriangleMesh:
       height: float,
       texture_type: Optional[mujoco.mjtTexture],
       resolution: int
-  ) -> TriangleMesh:
+  ) -> "TriangleMesh":
     """Creates a cylinder."""
     vertices = []
     triangles = []
diff --git a/python/mujoco/usd/utils.py b/python/mujoco/usd/utils.py
index 1e42afa2..ce3402ff 100644
--- a/python/mujoco/usd/utils.py
+++ b/python/mujoco/usd/utils.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Utility functions for USD exporter."""
 import numpy as np