make2/tools/optimize_model/optimize_model.py

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)