Optimized numpy processing when converting STL, OBJ, and DAE (#54)

Author: ft-lab

urdf_usd_converter/_impl/conversion_collada.py

@@ -62,11 +62,13 @@ def _convert_mesh(
     matrix = _multiply_root_matrix(_collada, matrix)
 
     all_face_vertex_counts: list[int] = []
-    all_face_vertex_indices: list[int] = []
+    all_face_vertex_indices_list: list[np.ndarray] = []
+    all_normals_list: list[np.ndarray] = []
     all_normals: Vt.Vec3fArray | None = None
-    all_normal_indices: list[int] = []
+    all_normal_indices_list: list[np.ndarray] = []
+    all_uvs_list: list[np.ndarray] = []
     all_uvs: Vt.Vec2fArray | None = None
-    all_uv_indices: list[int] = []
+    all_uv_indices_list: list[np.ndarray] = []
     face_offsets: list[int] = []
     face_material_names: list[str] = []
     current_normal_offset = 0
@@ -90,15 +92,16 @@ def _convert_mesh(
         # vertex indices.
         if is_triangle_type:
             face_vertex_counts, face_vertex_indices = convert_face_indices_array(primitive.vertex_index)
+            face_vertex_indices_array = np.array(face_vertex_indices, dtype=np.int32)
         else:  # Polylist or Polygons
             # Use numpy for faster conversion
             face_vertex_counts = primitive.vcounts.tolist()
-            face_vertex_indices = primitive.vertex_index.tolist()
+            face_vertex_indices_array = primitive.vertex_index
         all_face_vertex_counts.extend(face_vertex_counts)
-        all_face_vertex_indices.extend(face_vertex_indices)
+        all_face_vertex_indices_list.append(face_vertex_indices_array)
 
         # Remove duplicates and add used vertex indices.
-        unique_vertex_indices.extend(np.unique(face_vertex_indices))
+        unique_vertex_indices.extend(np.unique(face_vertex_indices_array))
 
         face_offsets.append(len(face_vertex_counts))
 
@@ -110,49 +113,56 @@ def _convert_mesh(
 
         # normals.
         if hasattr(primitive, "normal") and len(primitive.normal) > 0:
-            primitive_normals = convert_vec3f_array(primitive.normal)
-            all_normals = primitive_normals if all_normals is None else Vt.Vec3fArray(list(all_normals) + list(primitive_normals))
+            primitive_normals = np.array(primitive.normal, dtype=np.float32).reshape(-1, 3)
+            all_normals_list.append(primitive_normals)
             normal_indices = primitive.normal_index
 
             # Optimize flattening operation using numpy when possible
             if is_triangle_type:
                 # Flatten 2D array more efficiently
                 if isinstance(normal_indices, np.ndarray):
-                    normal_indices = normal_indices.ravel().tolist()
+                    normal_indices = normal_indices.ravel()
             else:  # Polylist or Polygons
-                normal_indices = normal_indices.tolist()
+                pass  # normal_indices is already a numpy array
 
-            normal_indices = (np.array(normal_indices, dtype=np.int32) + current_normal_offset).tolist()
-            all_normal_indices.extend(normal_indices)
+            normal_indices_array = np.array(normal_indices, dtype=np.int32) + current_normal_offset
+            all_normal_indices_list.append(normal_indices_array)
             current_normal_offset += len(primitive_normals)
 
         # uvs.
         if hasattr(primitive, "texcoordset") and len(primitive.texcoordset) > 0:
-            uv_data = convert_vec2f_array(primitive.texcoordset[0])
-            all_uvs = uv_data if all_uvs is None else Vt.Vec2fArray(list(all_uvs) + list(uv_data))
-            uv_indices = primitive.texcoord_index if hasattr(primitive, "texcoord_index") else list(range(len(uv_data)))
-            uv_indices = (np.array(uv_indices, dtype=np.int32) + current_uv_offset).tolist()
-            all_uv_indices.extend(uv_indices)
+            uv_data = np.array(primitive.texcoordset[0], dtype=np.float32).reshape(-1, 2)
+            all_uvs_list.append(uv_data)
+            uv_indices = primitive.texcoord_index if hasattr(primitive, "texcoord_index") else np.arange(len(uv_data), dtype=np.int32)
+            uv_indices_array = np.array(uv_indices, dtype=np.int32) + current_uv_offset
+            all_uv_indices_list.append(uv_indices_array)
             current_uv_offset += len(uv_data)
 
+    # Concatenate all numpy arrays into single arrays
+    all_face_vertex_indices = np.concatenate(all_face_vertex_indices_list) if all_face_vertex_indices_list else np.array([], dtype=np.int32)
+    all_normal_indices = np.concatenate(all_normal_indices_list) if all_normal_indices_list else np.array([], dtype=np.int32)
+    all_uv_indices = np.concatenate(all_uv_indices_list) if all_uv_indices_list else np.array([], dtype=np.int32)
+
     if len(all_face_vertex_counts) > 0 and len(all_face_vertex_indices) > 0 and all_vertices is not None:
         # Remove unused vertices from all_vertices and update the vertex list all_face_vertex_indices.
         unique_vertex_indices = np.unique(unique_vertex_indices)
 
         vertices_array = np.array(all_vertices, dtype=np.float32).reshape(-1, 3)
         all_vertices = convert_vec3f_array(vertices_array[unique_vertex_indices])
-        all_face_vertex_indices = np.searchsorted(unique_vertex_indices, all_face_vertex_indices).tolist()
+        all_face_vertex_indices = Vt.IntArray.FromNumpy(np.searchsorted(unique_vertex_indices, all_face_vertex_indices))
 
         # create a normal primvar data for the geometry.
         normals = None
-        if all_normals and all_normal_indices and len(all_normal_indices) == len(all_face_vertex_indices):
-            normals = usdex.core.Vec3fPrimvarData(UsdGeom.Tokens.faceVarying, all_normals, Vt.IntArray(all_normal_indices))
+        all_normals = convert_vec3f_array(np.concatenate(all_normals_list)) if len(all_normals_list) > 0 else None
+        if all_normals and len(all_normal_indices) > 0 and len(all_normal_indices) == len(all_face_vertex_indices):
+            normals = usdex.core.Vec3fPrimvarData(UsdGeom.Tokens.faceVarying, all_normals, Vt.IntArray.FromNumpy(all_normal_indices))
             normals.index()  # re-index the normals to remove duplicates
 
         # create a uv primvar data for the geometry.
         uvs = None
-        if all_uvs and all_uv_indices and len(all_uv_indices) == len(all_face_vertex_indices):
-            uvs = usdex.core.Vec2fPrimvarData(UsdGeom.Tokens.faceVarying, all_uvs, Vt.IntArray(all_uv_indices))
+        all_uvs = convert_vec2f_array(np.concatenate(all_uvs_list)) if len(all_uvs_list) > 0 else None
+        if all_uvs and len(all_uv_indices) > 0 and len(all_uv_indices) == len(all_face_vertex_indices):
+            uvs = usdex.core.Vec2fPrimvarData(UsdGeom.Tokens.faceVarying, all_uvs, Vt.IntArray.FromNumpy(all_uv_indices))
             uvs.index()  # re-index the uvs to remove duplicates
 
         # If only one geometry exists within the dae, only one mesh will be placed.
@@ -172,8 +182,8 @@ def _convert_mesh(
             _prim,
             _safe_name,
             faceVertexCounts=Vt.IntArray(all_face_vertex_counts),
-            faceVertexIndices=Vt.IntArray(all_face_vertex_indices),
-            points=Vt.Vec3fArray(all_vertices),
+            faceVertexIndices=all_face_vertex_indices,
+            points=all_vertices,
             normals=normals,
             uvs=uvs,
         )

urdf_usd_converter/_impl/mesh.py

@@ -134,16 +134,14 @@ def _convert_single_obj(
     points = convert_vec3f_array(np.asarray(vertices, dtype=np.float32).reshape(-1, 3))
 
     normals = None
-    source_normals = attrib.normals
-    if len(source_normals) > 0:
-        normals_data = convert_vec3f_array(np.asarray(source_normals, dtype=np.float32).reshape(-1, 3))
+    if len(attrib.normals) > 0:
+        normals_data = convert_vec3f_array(np.asarray(attrib.normals, dtype=np.float32).reshape(-1, 3))
         normals = usdex.core.Vec3fPrimvarData(UsdGeom.Tokens.faceVarying, normals_data, Vt.IntArray(obj_mesh.normal_indices()))
         normals.index()  # re-index the normals to remove duplicates
 
     uvs = None
-    source_uvs = attrib.texcoords
-    if len(source_uvs) > 0:
-        uv_data = convert_vec2f_array(np.asarray(source_uvs, dtype=np.float32).reshape(-1, 2))
+    if len(attrib.texcoords) > 0:
+        uv_data = convert_vec2f_array(np.asarray(attrib.texcoords, dtype=np.float32).reshape(-1, 2))
         uvs = usdex.core.Vec2fPrimvarData(UsdGeom.Tokens.faceVarying, uv_data, Vt.IntArray(obj_mesh.texcoord_indices()))
         uvs.index()  # re-index the uvs to remove duplicates
 
@@ -195,7 +193,7 @@ def convert_obj(prim: Usd.Prim, input_path: pathlib.Path, data: ConversionData)
 
     for shape, name, safe_name in zip(shapes, names, safe_names):
         obj_mesh = shape.mesh
-        face_vertex_counts = obj_mesh.num_face_vertices
+        face_vertex_counts = Vt.IntArray(obj_mesh.num_face_vertices)
         material_ids = obj_mesh.material_ids[0]
         material = materials[material_ids] if material_ids >= 0 else None
 
@@ -211,7 +209,7 @@ def convert_obj(prim: Usd.Prim, input_path: pathlib.Path, data: ConversionData)
         points = convert_vec3f_array(np.asarray(points_array, dtype=np.float32).reshape(-1, 3))
 
         # Remap indices using NumPy searchsorted
-        face_vertex_indices = np.searchsorted(unique_vertex_indices, vertex_indices_in_shape).tolist()
+        face_vertex_indices = Vt.IntArray.FromNumpy(np.searchsorted(unique_vertex_indices, vertex_indices_in_shape))
 
         # Process normals
         normals = None
@@ -220,11 +218,10 @@ def convert_obj(prim: Usd.Prim, input_path: pathlib.Path, data: ConversionData)
             unique_normal_indices = np.unique(normal_indices_in_shape)
 
             normals_array = np.array(attrib.normals, dtype=np.float32).reshape(-1, 3)
-            normals_data_array = normals_array[unique_normal_indices]
-            normals_data = convert_vec3f_array(np.asarray(normals_data_array, dtype=np.float32).reshape(-1, 3))
+            normals_data = convert_vec3f_array(normals_array[unique_normal_indices])
 
-            remapped_normal_indices = np.searchsorted(unique_normal_indices, normal_indices_in_shape).tolist()
-            normals = usdex.core.Vec3fPrimvarData(UsdGeom.Tokens.faceVarying, normals_data, Vt.IntArray(remapped_normal_indices))
+            remapped_normal_indices = Vt.IntArray.FromNumpy(np.searchsorted(unique_normal_indices, normal_indices_in_shape))
+            normals = usdex.core.Vec3fPrimvarData(UsdGeom.Tokens.faceVarying, normals_data, remapped_normal_indices)
             normals.index()  # re-index the normals to remove duplicates
 
         # Process UV coordinates
@@ -234,18 +231,17 @@ def convert_obj(prim: Usd.Prim, input_path: pathlib.Path, data: ConversionData)
             unique_texcoord_indices = np.unique(texcoord_indices_in_shape)
 
             texcoords_array = np.array(attrib.texcoords, dtype=np.float32).reshape(-1, 2)
-            uv_data_array = texcoords_array[unique_texcoord_indices]
-            uv_data = convert_vec2f_array(np.asarray(uv_data_array, dtype=np.float32).reshape(-1, 2))
+            uv_data = convert_vec2f_array(texcoords_array[unique_texcoord_indices])
 
-            remapped_texcoord_indices = np.searchsorted(unique_texcoord_indices, texcoord_indices_in_shape).tolist()
-            uvs = usdex.core.Vec2fPrimvarData(UsdGeom.Tokens.faceVarying, uv_data, Vt.IntArray(remapped_texcoord_indices))
+            remapped_texcoord_indices = Vt.IntArray.FromNumpy(np.searchsorted(unique_texcoord_indices, texcoord_indices_in_shape))
+            uvs = usdex.core.Vec2fPrimvarData(UsdGeom.Tokens.faceVarying, uv_data, remapped_texcoord_indices)
             uvs.index()  # re-index the uvs to remove duplicates
 
         usd_mesh = usdex.core.definePolyMesh(
             prim,
             safe_name,
-            faceVertexCounts=Vt.IntArray(face_vertex_counts),
-            faceVertexIndices=Vt.IntArray(face_vertex_indices),
+            faceVertexCounts=face_vertex_counts,
+            faceVertexIndices=face_vertex_indices,
             points=points,
             normals=normals,
             uvs=uvs,

urdf_usd_converter/_impl/numpy.py

@@ -20,13 +20,12 @@ def convert_vec2f_array(source: np.ndarray) -> Vt.Vec2fArray:
     Raises:
         AssertionError: If the second dimension of the input array is not divisible by 2.
     """
-    num_elements, element_size = source.shape
+    _, element_size = source.shape
     assert element_size % 2 == 0
 
     # Reshape to (total_vectors, 2) and create Vec2f objects in batch
     reshaped = source.reshape(-1, 2).astype(np.float32) if element_size != 2 else source
-    result = [Gf.Vec2f(float(v[0]), float(v[1])) for v in reshaped]
-    return Vt.Vec2fArray(result)
+    return Vt.Vec2fArray.FromNumpy(reshaped)
 
 
 def convert_vec3f_array(source: np.ndarray) -> Vt.Vec3fArray:
@@ -43,16 +42,15 @@ def convert_vec3f_array(source: np.ndarray) -> Vt.Vec3fArray:
     Raises:
         AssertionError: If the second dimension of the input array is not divisible by 3.
     """
-    num_elements, element_size = source.shape
+    _, element_size = source.shape
     assert element_size % 3 == 0
 
     # In the case of STL, element_size=9, and it holds the three vertices of a triangle as a one-dimensional array.
     # Therefore, we need to reshape the array to (total_vectors, 3) before creating the Vec3f objects.
 
     # Reshape to (total_vectors, 3) and create Vec3f objects in batch
     reshaped = source.reshape(-1, 3).astype(np.float32) if element_size != 3 else source
-    result = [Gf.Vec3f(float(v[0]), float(v[1]), float(v[2])) for v in reshaped]
-    return Vt.Vec3fArray(result)
+    return Vt.Vec3fArray.FromNumpy(reshaped)
 
 
 def convert_face_indices_array(source: np.ndarray) -> tuple[list[int], list[int]]: