From d7447c4a842dedf95b631ed53f3bbf448773c053 Mon Sep 17 00:00:00 2001
From: Abhishek Joshi <abhishekjoshinj@gmail.com>
Date: Sun, 21 Jul 2024 17:00:16 -0500
Subject: [PATCH] Partial cube mapping implementation in USD

---
 python/mujoco/usd/exporter.py |  2 +-
 python/mujoco/usd/objects.py  | 21 ++++++++------
 python/mujoco/usd/shapes.py   | 53 ++++++++++++++++++++++++++++++-----
 3 files changed, 59 insertions(+), 17 deletions(-)

diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py
index 832bce32..daa1fae4 100644
--- a/python/mujoco/usd/exporter.py
+++ b/python/mujoco/usd/exporter.py
@@ -209,7 +209,7 @@ class USDExporter:
     """Load textures."""
     data_adr = 0
     self.texture_files = []
-    for texture_id in tqdm.tqdm(range(self.model.ntex)):
+    for texture_id in range(self.model.ntex):
       texture_height = self.model.tex_height[texture_id]
       texture_width = self.model.tex_width[texture_id]
       texture_nchannel = self.model.tex_nchannel[texture_id]
diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py
index 1d68ccbb..79d8ac3f 100644
--- a/python/mujoco/usd/objects.py
+++ b/python/mujoco/usd/objects.py
@@ -363,19 +363,22 @@ class USDPrimitiveMesh(USDObject):
   def _get_uv_geometry(self):
     assert self.prim_mesh and self.prim_mesh.triangle_uvs is not None
 
-    s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid]
-
     mesh_texcoord = np.array(self.prim_mesh.triangle_uvs)
     mesh_facetexcoord = np.asarray(self.prim_mesh.triangles)
     
-    if self.model.mat_texuniform[self.geom.matid]:
-      if self.geom.size[0] > 0:
-        s_scale *= self.geom.size[0]
-      if self.geom.size[1] > 0:
-        t_scale *= self.geom.size[1]
+    geom_rgb_texture = self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB][1]
 
-    mesh_texcoord[:, 0] *= s_scale / (self.geom.size[0] * 2)
-    mesh_texcoord[:, 1] *= t_scale / (self.geom.size[1] * 2)
+    if geom_rgb_texture == mujoco.mjtTexture.mjTEXTURE_2D:
+      s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid]
+
+      if self.model.mat_texuniform[self.geom.matid]:
+        if self.geom.size[0] > 0:
+          s_scale *= self.geom.size[0]
+        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)
 
     return mesh_texcoord, mesh_facetexcoord.flatten()
 
diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py
index bc82b9a3..28b340f8 100644
--- a/python/mujoco/usd/shapes.py
+++ b/python/mujoco/usd/shapes.py
@@ -27,11 +27,48 @@ def get_triangle_uvs(
   if texture_type == None:
     return None
 
-  # (jabhi) assuming all mappings are 2d mapping temporarily
-  triangle_uvs = np.array([
-    [vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)]
-  )
-  return triangle_uvs
+  triangle_uvs = []
+  if texture_type == mujoco.mjtTexture.mjTEXTURE_2D:
+    triangle_uvs = [[vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)]
+    
+  elif texture_type == mujoco.mjtTexture.mjTEXTURE_CUBE:
+    center = np.mean(vertices, axis=0)
+    for vertex_id in np.nditer(triangles):
+      x, y, z = vertices[vertex_id] - center
+
+      abs_x, abs_y, abs_z = abs(x), abs(y), abs(z)
+      
+      if x > 0 and abs_x >= abs_y and abs_x >= abs_z:
+        u = -z / abs_x
+        v = y / abs_x
+      elif x <= 0 and abs_x >= abs_y and abs_x >= abs_z:
+        u = z / abs_x
+        v = y / abs_x
+      elif y > 0 and abs_y >= abs_x and abs_y >= abs_z:
+        u = x / abs_y
+        v = -z / abs_y
+      elif y <= 0 and abs_y >= abs_x and abs_y >= abs_z:
+        u = x / abs_y
+        v = z / abs_y
+      elif z > 0 and abs_z >= abs_x and abs_z >= abs_y:
+        u = x / abs_z
+        v = y / abs_z
+      elif z <= 0 and abs_z >= abs_x and abs_z >= abs_y:
+        u = -x / abs_z
+        v = y / abs_z
+
+      u = (u + 1.0) / 2.0
+      v = (v + 1.0) / 2.0
+      v /= 6
+
+      assert 0 <= u and u <= 1 and 0 <= v and v <= 1
+
+      triangle_uvs.append([u, v])
+  elif texture_type == mujoco.mjtTexture.mjTEXTURE_SKYBOX:
+    # defaults to 2D mapping temporarily
+    triangle_uvs = [[vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)]
+
+  return np.array(triangle_uvs)
 
 class TriangleMesh():
   """ Store UV and geometry information for a primitve mesh
@@ -210,7 +247,9 @@ class TriangleMesh():
       new_vertices = np.vstack((self.vertices, other.vertices))
       other_triangles = other.triangles + len(self.vertices)
       new_triangles = np.vstack((self.triangles, other_triangles))
-      new_triangle_uvs = np.vstack((self.triangle_uvs, other.triangle_uvs))
+      new_triangle_uvs = None
+      if self.triangle_uvs is not None:
+        new_triangle_uvs = np.vstack((self.triangle_uvs, other.triangle_uvs))
       return TriangleMesh(new_vertices, new_triangles, new_triangle_uvs)
     raise TypeError(f"Cannot add TriangleMesh with {type(other)}")
 
@@ -304,7 +343,7 @@ def mesh_config_generator(
 def mesh_factory(
     mesh_config: Dict[str, Any],
     texture_type: Optional[mujoco.mjtTexture],
-    resolution: int = 10,
+    resolution: int = 100,
 ):
   """Generates a mesh given a config consisting of shapes."""
   assert "name" in mesh_config