Spaces:
Build error
Build error
| # Modified from https://github.com/buaacyw/MeshAnything | |
| import mesh2sdf.core | |
| import numpy as np | |
| import skimage.measure | |
| import trimesh | |
| import time | |
| from typing import List, Tuple | |
| class MeshProcessor: | |
| """A class to handle mesh normalization, watertight conversion and point cloud sampling.""" | |
| def normalize_mesh_vertices(vertices: np.ndarray, scaling_factor: float = 0.95) -> Tuple[np.ndarray, np.ndarray, float]: | |
| """ | |
| Normalize mesh vertices to be centered at origin and scaled appropriately. | |
| """ | |
| min_bounds = vertices.min(axis=0) | |
| max_bounds = vertices.max(axis=0) | |
| center = (min_bounds + max_bounds) * 0.5 | |
| max_dimension = (max_bounds - min_bounds).max() | |
| scale = 2.0 * scaling_factor / max_dimension | |
| normalized_vertices = (vertices - center) * scale | |
| return normalized_vertices, center, scale | |
| def convert_to_watertight(mesh: trimesh.Trimesh, octree_depth: int = 7) -> trimesh.Trimesh: | |
| """ | |
| Convert to watertight using mesh2sdf and marching cubes. | |
| """ | |
| grid_size = 2 ** octree_depth | |
| iso_level = 2 / grid_size | |
| # Normalize vertices for SDF computation | |
| normalized_vertices, original_center, original_scale = MeshProcessor.normalize_mesh_vertices(mesh.vertices) | |
| # Compute signed distance field | |
| sdf = mesh2sdf.core.compute(normalized_vertices, mesh.faces, size=grid_size) | |
| # Run marching cubes algorithm | |
| vertices, faces, normals, _ = skimage.measure.marching_cubes(np.abs(sdf), iso_level) | |
| # Transform vertices back to original coordinate system | |
| vertices = vertices / grid_size * 2 - 1 # Map to [-1, 1] range | |
| vertices = vertices / original_scale + original_center | |
| # Create new watertight mesh | |
| watertight_mesh = trimesh.Trimesh(vertices, faces, normals=normals) | |
| return watertight_mesh | |
| def convert_meshes_to_point_clouds( | |
| meshes: List[trimesh.Trimesh], | |
| points_per_mesh: int = 8192, | |
| apply_marching_cubes: bool = False, | |
| octree_depth: int = 7 | |
| ) -> Tuple[List[np.ndarray], List[trimesh.Trimesh]]: | |
| """ | |
| Process a list of meshes into point clouds with normals. | |
| """ | |
| point_clouds_with_normals = [] | |
| processed_meshes = [] | |
| for mesh in meshes: | |
| # Optionally convert to watertight mesh | |
| if apply_marching_cubes: | |
| start_time = time.time() | |
| mesh = MeshProcessor.convert_to_watertight(mesh, octree_depth=octree_depth) | |
| processing_time = time.time() - start_time | |
| print(f"Marching cubes complete! Time: {processing_time:.2f}s") | |
| # Store processed mesh | |
| processed_meshes.append(mesh) | |
| # Sample points and get corresponding face normals | |
| points, face_indices = mesh.sample(points_per_mesh, return_index=True) | |
| point_normals = mesh.face_normals[face_indices] | |
| # Combine points and normals | |
| points_with_normals = np.concatenate([points, point_normals], axis=-1, dtype=np.float16) | |
| point_clouds_with_normals.append(points_with_normals) | |
| return point_clouds_with_normals |