Given a 3D point cloud, get unstructured mesh using Ball Pivoting algorithm

ball_pivoting.py

import open3d as o3d
import trimesh
import numpy as np


def ball_pivoting_reconstruction(xyz, radii=None):
    """Given a 3D point cloud, get unstructured mesh using ball pivoting algorithm
    
    Based on this stack overflow code snippet:
    https://stackoverflow.com/questions/56965268/how-do-i-convert-a-3d-point-cloud-ply-into-a-mesh-with-faces-and-vertices
    
    Parameters
    ----------
    xyz: [n_points, 3]
        input point cloud, numpy array
    radii: [n_radii]
        list of radiuses to use for the ball pivoting algorithm
    
    Returns
    -------
    mesh: trimesh Mesh object with verts, faces and normals of the mesh
    
    """
    
    # estimate normals first
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(xyz)
    pcd.estimate_normals()

    # heuristic to estimate the radius of a rolling ball
    if radii is None:
        distances = pcd.compute_nearest_neighbor_distance()
        avg_dist = np.mean(distances)
        radius = 1.5 * avg_dist   
        radii = [radius, radius * 2]
    
    mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_ball_pivoting(pcd, o3d.utility.DoubleVector(radii))

    mesh = trimesh.Trimesh(np.asarray(mesh.vertices), np.asarray(mesh.triangles), vertex_normals=np.asarray(mesh.vertex_normals))
    
    # try to fix normals with Trimesh
    mesh.fix_normals()
    
    # save mesh:
    # mesh.export('../logs/mesh.obj')
    
    return mesh