optimized meshdata transmit.

Author: aCodeDog

genesis/sensors/lidar/lidar.py

@@ -22,47 +22,6 @@
 from ..base_sensor import Sensor
 
 
-def get_link_mesh(link, use_visual_mesh=False):
-    """
-    Extract mesh from a given link.
-    
-    Args:
-        link: Genesis link object
-        use_visual_mesh: Whether to use visual mesh (True) or collision mesh (False)
-        
-    Returns:
-        trimesh.Trimesh: Combined mesh for the link
-    """
-    meshes = []
-    if use_visual_mesh:
-        geoms = link.vgeoms
-    else:
-        geoms = link.geoms
-        
-    for i, geom in enumerate(geoms):
-        T = trans_quat_to_T(
-            (geom.get_pos() - link.get_pos()).cpu(),
-            R_to_quat(
-                (
-                    quat_to_R(geom.get_quat())
-                    @ torch.linalg.inv(quat_to_R(link.get_quat()))
-                )
-            ).cpu(),
-        )
-        if T.ndim == 3:
-            T = T[0]  # NOTE: we use the canonical space so batch can be ignored
-        mesh = geom.get_trimesh().copy()  # NOTE: avoid in-place write
-        mesh.apply_transform(T)
-        meshes.append(mesh)
-        
-    if len(meshes) == 0:
-        # Return empty mesh if no geometry
-        return trimesh.Trimesh()
-        
-    combined_mesh = trimesh.util.concatenate(meshes)
-    return combined_mesh
-
-
 # Global constants for LiDAR ray casting - will be initialized when needed
 NO_HIT_RAY_VAL = None
 NO_HIT_SEGMENTATION_VAL = None
@@ -105,9 +64,6 @@ def __init__(self, vertices: np.ndarray, triangles: np.ndarray):
         self.vertices.from_numpy(vertices.astype(np.float32))
         self.triangles.from_numpy(triangles.astype(np.int32))
 
-        # Pre-store numpy arrays for kernel usage (static mesh optimization)
-        self.vertices_np = vertices.astype(np.float32)
-        self.triangles_np = triangles.astype(np.int32)
 
         # Build BVH acceleration structure
         self._build_bvh()
@@ -265,9 +221,9 @@ def ray_aabb_intersection(ray_start, ray_dir, aabb_min, aabb_max):
 
 @ti.kernel 
 def lidar_cast_rays_kernel_bvh(
-    # Mesh data
-    mesh_vertices: ti.types.ndarray(ndim=2),     # [n_vertices, 3]
-    mesh_triangles: ti.types.ndarray(ndim=2),    # [n_triangles, 3]
+    # Mesh data (now using Taichi fields directly - no CPU->GPU transfer!)
+    mesh_vertices: ti.template(),                # GPU Taichi field [n_vertices, 3]
+    mesh_triangles: ti.template(),               # GPU Taichi field [n_triangles, 3]
 
     # BVH data structures
     bvh_nodes: ti.template(),                    # The BVH node tree
@@ -285,7 +241,7 @@ def lidar_cast_rays_kernel_bvh(
     world_frame: ti.i32,
 ):
     """
-    Taichi kernel for LiDAR ray casting, accelerated by a Bounding Volume Hierarchy (BVH).
+    Taichi kernel for LiDAR ray casting.
     """
     n_triangles = mesh_triangles.shape[0]
 
@@ -342,15 +298,11 @@ def lidar_cast_rays_kernel_bvh(
                     # We need to find the original triangle index.
                     sorted_leaf_idx = node_idx - (n_triangles - 1)
                     original_tri_idx = bvh_morton_codes[0, sorted_leaf_idx][1]
+                    tri_indices = mesh_triangles[original_tri_idx]
+                    v0 = mesh_vertices[tri_indices[0]]
+                    v1 = mesh_vertices[tri_indices[1]]
+                    v2 = mesh_vertices[tri_indices[2]]
 
-                    # Get triangle vertices
-                    v0_idx = mesh_triangles[original_tri_idx, 0]
-                    v1_idx = mesh_triangles[original_tri_idx, 1]
-                    v2_idx = mesh_triangles[original_tri_idx, 2]
-                    
-                    v0 = ti.math.vec3(mesh_vertices[v0_idx, 0], mesh_vertices[v0_idx, 1], mesh_vertices[v0_idx, 2])
-                    v1 = ti.math.vec3(mesh_vertices[v1_idx, 0], mesh_vertices[v1_idx, 1], mesh_vertices[v1_idx, 2])
-                    v2 = ti.math.vec3(mesh_vertices[v2_idx, 0], mesh_vertices[v2_idx, 1], mesh_vertices[v2_idx, 2])
 
                     # Perform the expensive ray-triangle intersection test
                     hit_result = ray_triangle_intersection(lidar_position, ray_direction_world, v0, v1, v2)
@@ -417,8 +369,8 @@ def cast_rays(self, lidar_positions, lidar_quaternions, ray_vectors, far_plane,
 
         # Call the BVH-accelerated kernel
         lidar_cast_rays_kernel_bvh(
-            self.mesh_data.vertices_np, 
-            self.mesh_data.triangles_np,
+            self.mesh_data.vertices, 
+            self.mesh_data.triangles,
             self.mesh_data.bvh.nodes,
             self.mesh_data.bvh.morton_codes,
             lidar_positions, lidar_quaternions, ray_vectors,