Add UVs to FEM solver

Author: alelievr

genesis/engine/entities/fem_entity.py

@@ -389,6 +389,7 @@ def sample(self):
         from genesis.engine.materials.FEM.cloth import Cloth as ClothMaterial
 
         is_cloth = isinstance(self.material, ClothMaterial)
+        self._uvs = None  # Initialize UVs storage (indexed same as FEM vertices)
 
         if is_cloth:
             # Cloth: load surface mesh directly (no tetrahedralization)
@@ -400,6 +401,14 @@ def sample(self):
                 faces = mesh.faces
                 # For cloth, we store faces as "elements" (treating them as surface elements)
                 self.instantiate(verts, faces)
+
+                # Load UVs from mesh (1:1 mapping for cloth)
+                try:
+                    uvs = mesh.visual.uv.copy()
+                    uvs[:, 1] = 1.0 - uvs[:, 1]  # flip V coordinate (trimesh uses top-left origin)
+                    self._uvs = uvs.astype(gs.np_float)
+                except AttributeError:
+                    self._uvs = None
             else:
                 gs.raise_exception(f"Cloth material only supports Mesh morph. Got: {self.morph}.")
         else:
@@ -419,11 +428,12 @@ def sample(self):
             elif isinstance(self.morph, gs.options.morphs.Cylinder):
                 verts, elems = eu.cylinder_to_elements()
             elif isinstance(self.morph, gs.options.morphs.Mesh):
-                verts, elems = eu.mesh_to_elements(
+                verts, elems, self._uvs = eu.mesh_to_elements(
                     file=self._morph.file,
                     pos=self._morph.pos,
                     scale=self._morph.scale,
                     tet_cfg=self.tet_cfg,
+                    return_uvs=True,
                 )
             else:
                 gs.raise_exception(f"Unsupported morph: {self.morph}.")
@@ -477,8 +487,23 @@ def _add_to_solver(self, in_backward=False):
                 tri2el=self._surface_el_np,
             )
 
+        # Add UVs to solver if available
+        self._add_uvs_to_solver()
+
         self.active = True
 
+    def _add_uvs_to_solver(self):
+        """
+        Copy UV coordinates to the solver's surface_render_uvs field.
+        """
+        if self._uvs is None:
+            return
+
+        # Copy UVs to solver field starting at v_start
+        uvs_np = self._uvs.astype(gs.np_float, copy=False)
+        for i, uv in enumerate(uvs_np):
+            self._solver.surface_render_uvs[self._v_start + i] = uv
+
     def compute_pressure_field(self):
         """
         Compute the pressure field for the FEM entity based on its tetrahedral elements.
@@ -1074,6 +1099,11 @@ def surface_triangles(self):
         """Surface triangles of the FEM mesh."""
         return self._surface_tri_np
 
+    @property
+    def uvs(self):
+        """UV coordinates for this entity's vertices, or None if not available."""
+        return self._uvs
+
     @property
     def tet_cfg(self):
         """Configuration of tetrahedralization."""

genesis/engine/solvers/fem_solver.py

@@ -240,6 +240,13 @@ def init_surface_fields(self):
             layout=ti.Layout.SOA,
         )
 
+        # UV coordinates for rendering (per-vertex UVs, initialized to zeros)
+        self.surface_render_uvs = ti.field(
+            dtype=gs.ti_vec2,
+            shape=(max(n_vertices_max, 1),),
+            needs_grad=False,
+        )
+
     def _init_surface_info(self):
         self.vertices_on_surface = ti.field(dtype=gs.ti_bool, shape=(self.n_vertices,))
         self.elements_on_surface = ti.field(dtype=gs.ti_bool, shape=(self.n_elements,))
@@ -1096,8 +1103,9 @@ def get_state_render(self, f):
         self.get_state_render_kernel(f)
         vertices = self.surface_render_v.vertices
         indices = self.surface_render_f.indices
+        uvs = self.surface_render_uvs
 
-        return vertices, indices
+        return vertices, indices, uvs
 
     def get_forces(self):
         """
@@ -1378,7 +1386,8 @@ def get_state_render_kernel(self, f: ti.i32):
                 pos_j = ti.cast(self.elements_v[f, i_v, i_b].pos[j], ti.f32)
                 self.surface_render_v[i_v, i_b].vertices[j] = pos_j + self.envs_offset[i_b][j]
 
-        for i_s, i_b in ti.ndrange(self.n_surfaces, self._B):
+        # Fill triangle indices (flat array, 3 ints per triangle)
+        for i_s in range(self.n_surfaces):
             for j in ti.static(range(3)):
                 self.surface_render_f[i_s * 3 + j].indices = ti.cast(self.surface[i_s].tri2v[j], ti.i32)
 

genesis/utils/element.py

@@ -31,9 +31,21 @@ def cylinder_to_elements():
     raise NotImplementedError
 
 
-def mesh_to_elements(file, pos=(0, 0, 0), scale=1.0, tet_cfg=dict()):
-    mesh = mu.load_mesh(file)
+def mesh_to_elements(file, pos=(0, 0, 0), scale=1.0, tet_cfg=dict(), return_uvs=False):
+    # Load mesh (with texture data if UVs are requested)
+    mesh = mu.load_mesh(file, skip_texture=not return_uvs)
     mesh.vertices = mesh.vertices * scale
+    
+    # Extract UVs from mesh before tetrahedralization (if requested)
+    # (tetgen preserves original vertices at start of output array)
+    if return_uvs:
+        n_original = len(mesh.vertices)
+        try:
+            original_uvs = mesh.visual.uv.copy()
+            original_uvs[:, 1] = 1.0 - original_uvs[:, 1]  # flip V (trimesh uses top-left origin)
+            original_uvs = original_uvs.astype(gs.np_float)
+        except AttributeError:
+            original_uvs = None
 
     # compute file name via hashing for caching
     tet_file_path = mu.get_tet_path(mesh.vertices, mesh.faces, tet_cfg)
@@ -43,8 +55,8 @@ def mesh_to_elements(file, pos=(0, 0, 0), scale=1.0, tet_cfg=dict()):
     if os.path.exists(tet_file_path):
         gs.logger.debug("Tetrahedra file (`.tet`) found in cache.")
         try:
-            with open(tet_file_path, "rb") as file:
-                verts, elems = pkl.load(file)
+            with open(tet_file_path, "rb") as tet_file:
+                verts, elems = pkl.load(tet_file)
             is_cached_loaded = True
         except (EOFError, ModuleNotFoundError, pkl.UnpicklingError, TypeError, MemoryError):
             gs.logger.info("Ignoring corrupted cache.")
@@ -54,10 +66,18 @@ def mesh_to_elements(file, pos=(0, 0, 0), scale=1.0, tet_cfg=dict()):
             verts, elems = mu.tetrahedralize_mesh(mesh, tet_cfg)
 
             os.makedirs(os.path.dirname(tet_file_path), exist_ok=True)
-            with open(tet_file_path, "wb") as file:
-                pkl.dump((verts, elems), file)
+            with open(tet_file_path, "wb") as tet_file:
+                pkl.dump((verts, elems), tet_file)
 
     verts += np.array(pos)
+    
+    if return_uvs:
+        # Build full UV array: original vertices get their UVs, interior vertices get zeros
+        n_fem = len(verts)
+        uvs = np.zeros((n_fem, 2), dtype=gs.np_float)
+        if original_uvs is not None:
+            uvs[:n_original] = original_uvs
+        return verts, elems, uvs
 
     return verts, elems
 

genesis/utils/mesh.py

@@ -169,9 +169,9 @@ def get_hashkey(*args):
     return hasher.hexdigest()
 
 
-def load_mesh(file):
+def load_mesh(file, skip_texture=True):
     if isinstance(file, (str, Path)):
-        return trimesh.load(file, force="mesh", skip_texture=True)
+        return trimesh.load(file, force="mesh", skip_texture=skip_texture)
     return file