新的信息

tools/optimize_model/optimize_model.py

@@ -0,0 +1,244 @@
+import bpy
+import sys
+import open3d as o3d
+import numpy as np
+from sklearn.neighbors import NearestNeighbors
+import os
+import argparse
+
+def estimate_point_curvatures(point_cloud, k=30):
+    point_cloud_points = np.asarray(point_cloud.points)
+    # Initialize the nearest neighbors finder
+    nn = NearestNeighbors(n_neighbors=k)
+    nn.fit(point_cloud_points)
+    curvatures = {i: 0 for i in range(len(point_cloud.points))}
+    for i, point in enumerate(point_cloud_points):
+        # Find the k-nearest neighbors (including the point itself)
+        distances, indices = nn.kneighbors([point])
+        neighbors = point_cloud_points[indices[0]]
+        # Compute the covariance matrix
+        covariance_matrix = np.cov(neighbors - neighbors.mean(axis=0), rowvar=False)
+        # Compute the eigenvalues (sorted)
+        eigenvalues, _ = np.linalg.eigh(covariance_matrix)
+        eigenvalues.sort()
+        # Calculate curvature using the smallest eigenvalue
+        curvature = eigenvalues[0] / sum(eigenvalues)
+        curvatures[i] = curvature
+    return curvatures
+
+def compute_adjacent_normal_angle_variance(pcd, pcd_normal, distance_scale):
+    kdtree = o3d.geometry.KDTreeFlann(pcd)
+    distances = []
+    for i in range(len(pcd.points)):
+        [k, idx, dist] = kdtree.search_knn_vector_3d(pcd.points[i], 2)
+        dist = dist[1:]  # distances to the actual neighbors
+        distances.append(dist)
+    mean_distance = np.mean(distances)
+    adjacent_normal_angles_variance = {i: 0 for i in range(len(np.asarray(pcd.points)))}
+    for i, point in enumerate(pcd.points):
+        # Search for the k-nearest neighbors of the point, up to 1500
+        k, idx, dist = kdtree.search_knn_vector_3d(point, 200)
+        idx = np.asarray(idx)
+        indices = idx[dist < mean_distance * distance_scale]
+        current_normal = pcd_normal[i]
+        # Calculate dot products between the current normal and its neighbors' normals
+        dot_products = np.dot(pcd_normal[indices], current_normal)
+        # Handle any potential NaN values in dot products
+        valid_dot_products = np.isfinite(dot_products)
+        dot_products = dot_products[valid_dot_products]
+        # Clip dot product values to valid range for arccos
+        dot_products = np.clip(dot_products, -1.0, 1.0)
+        # Compute angles in degrees and adjust for any NaNs which might be due to numerical errors
+        angles_array = np.degrees(np.arccos(dot_products))
+        # Set the angle with itself to zero since it's not meaningful
+        angles_array[indices == i] = 0
+        # Calculate the mean angle excluding the zero for the current point
+        valid_angles = angles_array[indices != i]
+        if valid_angles.size > 0:
+            adjacent_normal_angles_variance[i] = np.std(valid_angles)
+        else:
+            adjacent_normal_angles_variance[i] = 0  # Default to 0 if no valid neighbors exist
+    return adjacent_normal_angles_variance
+
+def compute_adjacent_point_vector_degrees_mean(pcd, pcd_normal, distance_scale):
+    kdtree = o3d.geometry.KDTreeFlann(pcd)
+    distances = []
+    for i in range(len(pcd.points)):
+        [k, idx, dist] = kdtree.search_knn_vector_3d(pcd.points[i], 2)
+        dist = dist[1:]  # distances to the actual neighbors
+        distances.append(dist)
+    mean_distance = np.mean(distances)
+    adjacent_point_vector_degrees_mean = {i: 0 for i in range(len(np.asarray(pcd.points)))}
+    for i, point in enumerate(pcd.points):
+        # Search for the k-nearest neighbors of the point, up to 200
+        k, idx, dist = kdtree.search_knn_vector_3d(point, 200)
+        idx = np.asarray(idx)
+        indices = idx[dist < mean_distance * distance_scale]
+        indices = indices[indices != i]
+        # Retrieve the normal of the current point
+        current_normal = pcd_normal[i]
+        adjacent_point_vector = np.asarray(pcd.points)[indices] - point
+        adjacent_point_vector_cos = np.dot(adjacent_point_vector, current_normal) / np.linalg.norm(adjacent_point_vector, axis = 1) * 1
+        valid_adjacent_point_vector_cos = np.isfinite(adjacent_point_vector_cos)
+        adjacent_point_vector_cos = adjacent_point_vector_cos[valid_adjacent_point_vector_cos]
+        if len(adjacent_point_vector_cos) > 0:
+            adjacent_point_vector_degrees = np.degrees(np.arccos(adjacent_point_vector_cos))
+            adjacent_point_vector_mean_degrees = np.mean(adjacent_point_vector_degrees)
+            adjacent_point_vector_degrees_mean[i] = adjacent_point_vector_mean_degrees
+        else:
+            adjacent_point_vector_degrees_mean[i] = 0
+    return adjacent_point_vector_degrees_mean
+
+def spread_point(pcd, selected_index, distance_scale, count = False):
+    kdtree = o3d.geometry.KDTreeFlann(pcd)
+    distances = []
+    for i in range(len(pcd.points)):
+        [k, idx, dist] = kdtree.search_knn_vector_3d(pcd.points[i], 2)
+        dist = dist[1:]  # distances to the actual neighbors
+        distances.append(dist)
+    mean_distance = np.mean(distances)
+    if count is False:
+        spread_selected_index = set()
+        for i in selected_index:
+            k, idx, dist = kdtree.search_knn_vector_3d(pcd.points[i], 200)
+            idx = np.asarray(idx)
+            dist = np.asarray(dist)
+            indices = idx[dist < mean_distance * distance_scale]
+            spread_selected_index.update(indices)
+        spread_selected_index = np.array(list(spread_selected_index))
+    else:
+        spread_selected_index = {i: 0 for i in range(len(np.asarray(pcd.points)))}
+        for i in selected_index:
+            k, idx, dist = kdtree.search_knn_vector_3d(pcd.points[i], 200)
+            idx = np.asarray(idx)
+            dist = np.asarray(dist)
+            indices = idx[dist < mean_distance * distance_scale]
+            for indice in indices:
+                spread_selected_index[indice] += 1
+    return spread_selected_index
+
+def ApplySmoothEffort(object, selected_vertice_indices = None, factor = 0.5, iterations = 30):
+    if selected_vertice_indices is not None:
+        # Assuming 'mesh' is already defined, e.g., mesh = bpy.context.object.data
+        selected_vertices_index = set()  # Initialize an empty set to store unique vertex indices
+        # Iterate over each index in the list of polygon indices
+        for idx in selected_vertice_indices:
+            polygon = object.data.polygons[idx]
+            for vert_idx in polygon.vertices:
+                selected_vertices_index.add(vert_idx)
+        # Switch to Edit mode
+        bpy.ops.object.mode_set(mode='EDIT') 
+        # Ensure we're dealing with the latest data
+        bpy.ops.object.mode_set(mode='OBJECT')
+        for i, vertex in enumerate(object.data.vertices):
+            vertex.select = False
+            if i in selected_vertices_index:
+                vertex.select = True
+    # Update the object to reflect changes
+    bpy.ops.object.mode_set(mode='EDIT')
+    bpy.ops.object.mode_set(mode='OBJECT')
+    # Add a Smooth modifier and set it to use the vertex group
+    smooth_mod = object.modifiers.new(name="SmoothMod", type='SMOOTH')
+    if selected_vertice_indices is not None:
+        # Create a new vertex group
+        group = object.vertex_groups.new(name='Smooth Group')
+        # Add selected vertices to the vertex group
+        group.add([v.index for v in object.data.vertices if v.select], weight=1.0, type='ADD')
+        smooth_mod.vertex_group = 'Smooth Group'
+    smooth_mod.factor = factor
+    smooth_mod.iterations = iterations
+    # bpy.ops.object.modifier_apply(modifier="SmoothMod")
+    # if selected_vertice_indices is not None:
+    #     object.vertex_groups.clear()
+
+def process(input_mesh_path, output_mesh_path):
+    # Select and delete all objects in the current scene to start fresh
+    bpy.ops.object.select_all(action='SELECT')
+    bpy.ops.object.delete(confirm=False)
+    # Path to the OBJ file to be imported
+    bpy.ops.wm.obj_import(filepath=input_mesh_path)
+    # Set the current object to the newly imported mesh
+    current_obj = bpy.context.object
+    # Retrieve the mesh data from the current object
+    mesh_data = current_obj.data
+    bpy.context.view_layer.update()
+    # Convert Blender mesh vertices to a NumPy array
+    vertices_array = np.array([vertex.co[:] for vertex in mesh_data.vertices])
+    # Convert Blender mesh polygons to a NumPy array of triangles
+    triangles_array = np.array([polygon.vertices[:] for polygon in mesh_data.polygons])
+
+    # Create an Open3D mesh object from the arrays
+    open3d_mesh = o3d.geometry.TriangleMesh()
+    open3d_mesh.vertices = o3d.utility.Vector3dVector(vertices_array)
+    open3d_mesh.triangles = o3d.utility.Vector3iVector(triangles_array)
+    open3d_mesh.compute_triangle_normals()
+    # Retrieve and store the normals as a NumPy array
+    triangle_normals_array = np.asarray(open3d_mesh.triangle_normals)
+
+    # Compute centroids for each face of the mesh
+    triangle_centroids_array = np.array([
+        np.mean(vertices_array[triangle], axis=0)
+        for triangle in open3d_mesh.triangles
+    ])
+    # Create a point cloud of face centroids
+    pcd_triangle_centroid = o3d.geometry.PointCloud()
+    pcd_triangle_centroid.points = o3d.utility.Vector3dVector(triangle_centroids_array)
+
+    point_curvatures = estimate_point_curvatures(pcd_triangle_centroid, 30)
+    np_point_curvatures = np.array(list(point_curvatures.values()))
+    point_curvatures_index = np.arange(len(np_point_curvatures))[np_point_curvatures > 0.25]
+
+    spread_point_curvatures_index = spread_point(pcd_triangle_centroid, point_curvatures_index, 50)
+
+    ApplySmoothEffort(current_obj, spread_point_curvatures_index, 0.5, 30)
+
+    adjacent_normal_angle_variance = compute_adjacent_normal_angle_variance(pcd_triangle_centroid, triangle_normals_array, 40)
+    adjacent_point_vector_degrees_mean = compute_adjacent_point_vector_degrees_mean(pcd_triangle_centroid, triangle_normals_array, 40)
+
+    array_adjacent_point_vector_degrees_mean = np.array(list(adjacent_point_vector_degrees_mean.values()))
+    array_adjacent_normal_angle_variance = np.array(list(adjacent_normal_angle_variance.values()))
+
+    use_data = np.where(array_adjacent_point_vector_degrees_mean < 90, array_adjacent_normal_angle_variance, 0)
+    # Calculate the minimum and maximum values in the angle array, ignoring NaNs.
+    min_angle = np.nanmin(use_data)
+    max_angle = np.nanmax(use_data)
+    # Perform min-max normalization on the angles to scale them between 0 and 1.
+    # This normalization helps in comparing values on a common scale.
+    normalized_angles = (use_data - min_angle) / (max_angle - min_angle)
+    first_select_thread = normalized_angles > 0.2
+    first_select_thread_number = first_select_thread.astype(np.float64)
+    thread_normalized_angles_index = np.arange(len(first_select_thread))[first_select_thread]
+
+    spread_point_count = spread_point(pcd_triangle_centroid, thread_normalized_angles_index, 3, count=True)
+
+    overlapp_trigular_points = np.array(list(spread_point_count.values())) > 3
+    overlapp_trigular_points_index = np.arange(len(overlapp_trigular_points))[overlapp_trigular_points]
+
+    ApplySmoothEffort(current_obj, overlapp_trigular_points_index, 0.5, 30)
+
+    ApplySmoothEffort(current_obj, factor=0.5, iterations=30)
+
+    bpy.ops.wm.obj_export(filepath=output_mesh_path)
+
+def process_meshes(input_path, output_path):
+    """
+    Processes each mesh file in the input directory and saves the cleaned meshes in the output directory.
+
+    Args:
+    input_path (str): Directory containing the input meshes or a single mesh file.
+    output_path (str): Directory where the cleaned meshes will be saved.
+    """
+    if os.path.isdir(input_path):
+        for filename in sorted(os.listdir(input_path)):
+            full_input_path = os.path.join(input_path, filename)
+            full_output_path = os.path.join(output_path, filename)
+            process(full_input_path, full_output_path)
+    else:
+        process(input_path, output_path)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Processes mesh files.")
+    parser.add_argument("-ip", "--input_path", type=str, required=True, help="Input path for the mesh file or directory.")
+    parser.add_argument("-op", "--output_path", type=str, required=True, help="Output path for the cleaned mesh file or directory.")
+    args = parser.parse_args()
+    process_meshes(args.input_path, args.output_path)