""" ----------------------------------------------------------------------------- Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. NVIDIA CORPORATION and its licensors retain all intellectual property and proprietary rights in and to this software, related documentation and any modifications thereto. Any use, reproduction, disclosure or distribution of this software and related documentation without an express license agreement from NVIDIA CORPORATION is strictly prohibited. ----------------------------------------------------------------------------- """ import sys sys.path.append(".") import argparse import glob import os import kiui import numpy as np import tqdm import trimesh from meshiki import Mesh parser = argparse.ArgumentParser() parser.add_argument("test_path", type=str, help="path to the mesh file or folder") parser.add_argument("--verbose", action="store_true", help="print verbose output") parser.add_argument("--force_cc", action="store_true", help="force to use connected components and ignore glb groups") parser.add_argument("--no_smart_group", action="store_true", help="do not perform smart grouping") parser.add_argument("--no_stitch", action="store_true", help="do not stitch open boundaries") parser.add_argument("--no_merge_odd_loops", action="store_true", help="do not merge odd loops") parser.add_argument("--no_dilate", action="store_true", help="do not dilate the mesh") parser.add_argument("--dilate_size", type=float, default=2 / 512, help="dilate size") parser.add_argument("--workspace", type=str, default="output", help="path to the output folder") opt = parser.parse_args() class NamedDisjointSet: def __init__(self, names): # names: list of str self.parent = {name: name for name in names} def find(self, x): if self.parent[x] != x: self.parent[x] = self.find(self.parent[x]) return self.parent[x] def merge(self, x, y): self.parent[self.find(x)] = self.find(y) def get_groups(self): groups = {} for name in self.parent: root = self.find(name) if root not in groups: groups[root] = [] groups[root].append(name) return groups def is_single_layer_plane(mesh: trimesh.Trimesh, coplane_thresh: float = 1): # check if a mesh is just a single-layer plane if mesh.is_watertight: return False face_normals = mesh.face_normals diff = np.linalg.norm(np.abs(face_normals) - np.abs(face_normals[0])) return diff.max() < coplane_thresh def calc_intersection_union(bounds1, bounds2): # bounds: [2, 3] bmin1, bmax1 = bounds1[0], bounds1[1] bmin2, bmax2 = bounds2[0], bounds2[1] # intersection bmin = np.maximum(bmin1, bmin2) bmax = np.minimum(bmax1, bmax2) if np.any(bmin >= bmax): intersection = 0 else: intersection = np.prod(bmax - bmin) # union vol1 = np.prod(bmax1 - bmin1) vol2 = np.prod(bmax2 - bmin2) union = vol1 + vol2 - intersection return intersection, union def is_coplanar_and_convex(vertices, coplanar_thresh: float = 1): # vertices: [N, 3], assume ordered to form a coplanar polygon # note the last vertex is the same as the first vertex, i.e. ABCA # we are actually not requiring perfect coplanarity, the thresh of 1 means 60 degree tolerance... # Need at least 3 vertices to form a polygon if len(vertices) < 3: return False # Convert to numpy array if not already points = np.array(vertices) n_points = len(points) # Normal of the first triangle v1 = points[1] - points[0] v2 = points[2] - points[0] normal = np.cross(v1, v2) normal = normal / (np.linalg.norm(normal) + 1e-12) # Test if coplanar using the normal of fan-cut triangles for i in range(2, n_points - 2): v1 = points[i] - points[0] v2 = points[i + 1] - points[0] normal_cur = np.cross(v1, v2) normal_cur = normal_cur / (np.linalg.norm(normal_cur) + 1e-12) diff = np.linalg.norm(np.abs(normal_cur) - np.abs(normal)) if diff > coplanar_thresh: # print(f'not coplanar: {normal_cur} != {normal} (diff = {diff:.4f}) at {i}-{i+1}') return False # not coplanar # Find basis vectors for the 2D plane # First basis vector can be the normalized vector from points[0] to points[1] basis1 = v1 / (np.linalg.norm(v1) + 1e-12) # Second basis vector is perpendicular to both normal and basis1 basis2 = np.cross(normal, basis1) basis2 = basis2 / (np.linalg.norm(basis2) + 1e-12) # Project all points onto the 2D plane points_2d = np.zeros((n_points, 2)) for i in range(n_points): v = points[i] - points[0] points_2d[i, 0] = np.dot(v, basis1) points_2d[i, 1] = np.dot(v, basis2) # Check if polygon is convex by using the cross product # For a convex polygon, all cross products should have the same sign sign = 0 for i in range(n_points - 1): j = (i + 1) % n_points k = (i + 2) % n_points # Vectors from point i to j and j to k v1 = points_2d[j] - points_2d[i] v2 = points_2d[k] - points_2d[j] # 2D cross product cross_product = v1[0] * v2[1] - v1[1] * v2[0] # Check for consistent sign of cross product if abs(cross_product) > 1e-2: # Skip collinear points current_sign = np.sign(cross_product) if sign == 0: sign = current_sign elif sign != current_sign: # print(f'not convex: {i}-{j} and {j}-{k}, cross_product = {cross_product:.4f}') return False # not convex return True def stitch_nonwatertight_mesh(mesh: trimesh.Trimesh, eps: float = 1e-2): # mesh will be inplace modified # return a flag denoting if there are still open boundaries unfixed # manager = trimesh.collision.CollisionManager() # manager.add_object('main', mesh) # watertight mesh doesn't need to be stitched if mesh.is_watertight: return False # planar mesh cannot be stitched if is_single_layer_plane(mesh, eps): return True # the following is modified from trimesh.repair.stitch # fan_faces = trimesh.repair.stitch(mesh) nonwatertight = False from trimesh.path.exchange.misc import faces_to_path faces = np.arange(len(mesh.faces)) # get a sequence of vertex indices representing the # boundary of the specified faces # will be referencing the same indexes of `mesh.vertices` boundaries = [ e.points for e in faces_to_path(mesh, faces)["entities"] if len(e.points) > 3 and e.points[0] == e.points[-1] ] # get properties to avoid querying in loop vertices = mesh.vertices normals = mesh.face_normals # find which faces are associated with an edge edges_face = mesh.edges_face tree_edge = mesh.edges_sorted_tree # MODIFIED: if any two boundary edges share close vertices, we discard both since they may connect mask = np.ones(len(boundaries), dtype=bool) for i in range(len(boundaries)): for j in range(i + 1, len(boundaries)): verts_i = vertices[boundaries[i]] # [N, 3] verts_j = vertices[boundaries[j]] # [M, 3] # check pair-wise distance dists = np.linalg.norm(verts_i[:, None, :] - verts_j[None, :, :], axis=-1) # [N, M] num_close = np.sum(dists < 1e-6) if num_close >= 4 or (num_close / verts_i.shape[0] >= 0.5) or (num_close / verts_j.shape[0] >= 0.5): mask[i] = False mask[j] = False # print(f'discarding boundary {i} and {j} because of close vertices') boundaries = [boundaries[i] for i in range(len(boundaries)) if mask[i]] # MODIFIED: we only keep coplanar & convex fans fans = [] for vert_indices in boundaries: # the fan should be coplanar and convex verts = vertices[vert_indices] if not is_coplanar_and_convex(verts): nonwatertight = True continue fan = np.column_stack( (np.ones(len(vert_indices) - 3, dtype=int) * vert_indices[0], vert_indices[1:-2], vert_indices[2:-1]) ) # [N, 3] fans.append(fan) # now we do a normal check against an adjacent face # to see if each region needs to be flipped for i, t in zip(range(len(fans)), fans): # get the edges from the original mesh # for the first `n` new triangles e = t[:10, 1:].copy() e.sort(axis=1) # find which indexes of `mesh.edges` these # new edges correspond with by finding edges # that exactly correspond with the tree query = tree_edge.query_ball_point(e, r=1e-10) if len(query) == 0: continue # stack all the indices that exist edge_index = np.concatenate(query) # get the normals from the original mesh original = normals[edges_face[edge_index]] # calculate the normals for a few new faces check, valid = trimesh.triangles.normals(vertices[t[:3]]) if not valid.any(): continue # take the first valid normal from our new faces check = check[0] # if our new faces are reversed from the original # Adjacent face flip them along their axis sign = np.dot(original, check) if sign.mean() < 0: fans[i] = np.fliplr(t) if len(fans) > 0: fans = np.vstack(fans) mesh.faces = np.concatenate([mesh.faces, fans]) return nonwatertight def smart_grouping(meshes: dict): # meshes: {name: trimesh.Trimesh, ...} # use collision manager to find all colliding pairs manager = trimesh.collision.CollisionManager() for name, mesh in meshes.items(): manager.add_object(name, mesh) is_collide, collide_pairs = manager.in_collision_internal(return_names=True) # print(f'[INFO] num_collide = {len(collide_pairs)}, {collide_pairs}') if not is_collide: return meshes # pre-calculate some stat for each mesh name_to_stat = {} total_volume = 0 max_extent = 0 num_submeshes = 0 num_meshes = len(meshes) for name, mesh in meshes.items(): name_to_stat[name] = {} submeshes = mesh.split() name_to_stat[name]["volume"] = [] name_to_stat[name]["extent"] = [] num_submeshes += len(submeshes) for submesh in submeshes: if submesh.is_watertight: name_to_stat[name]["volume"].append(submesh.volume) total_volume += name_to_stat[name]["volume"][-1] name_to_stat[name]["extent"].append(np.max(submesh.extents)) max_extent = max(max_extent, name_to_stat[name]["extent"][-1]) name_to_stat[name]["volume"] = ( np.mean(name_to_stat[name]["volume"]) if len(name_to_stat[name]["volume"]) > 0 else np.inf ) name_to_stat[name]["extent"] = ( np.max(name_to_stat[name]["extent"]) if len(name_to_stat[name]["extent"]) > 0 else np.inf ) # use a disjoint set to record grouping ds = NamedDisjointSet(list(meshes.keys())) # decide the merging thresh adaptively (based on the number of meshes, average volume and extent) # very empirical... if num_meshes <= 16: tol_volume = 0.05 * total_volume tol_extent = 0.05 * max_extent else: tol_volume = 0.1 * total_volume tol_extent = 0.1 * max_extent # loop each pair, determine if they should be grouped together for name1, name2 in collide_pairs: if ds.find(name1) == ds.find(name2): continue # single-layer plane should be merged if is_single_layer_plane(meshes[name1]) or is_single_layer_plane(meshes[name2]): # print(f'[INFO] merge {name1} and {name2} because of single-layer plane') ds.merge(name1, name2) continue # too small component if (name_to_stat[name1]["volume"] < tol_volume and name_to_stat[name1]["extent"] < tol_extent) or ( name_to_stat[name2]["volume"] < tol_volume and name_to_stat[name2]["extent"] < tol_extent ): # print(f'[INFO] merge {name1} and {name2} because of small volume') ds.merge(name1, name2) continue # overlaps a lot should be merged (just use bounding box IoU) bounds1 = meshes[name1].bounds # [2, 3] bounds2 = meshes[name2].bounds # [2, 3] vol_intersect, vol_union = calc_intersection_union(bounds1, bounds2) vol_iou = vol_intersect / vol_union if vol_iou > 0.5: # print(f'[INFO] merge {name1} and {name2} because of large IoU') ds.merge(name1, name2) continue # merge groups groups = ds.get_groups() for group in groups.values(): if len(group) <= 1: continue # print(f'[INFO] merge group: {group}') new_name = "_".join(group) new_mesh = trimesh.util.concatenate(list(meshes[name] for name in group)) # merge close vertices, and clean up new_mesh.merge_vertices(merge_tex=True, merge_norm=True) new_mesh.update_faces(new_mesh.unique_faces() & new_mesh.nondegenerate_faces()) new_mesh.fix_normals() meshes[new_name] = new_mesh for name in group: del meshes[name] print(f"[INFO] after grouping, num_meshes = {len(meshes)}") return meshes def merge_odd_loops(meshes: dict, graph: dict, penetration_depths: dict): # meshes: {name: trimesh.Trimesh, ...} # graph: {name: [neighbor, ...], ...} # penetration_depths: {(name1, name2): depth, ...} # find out odd loops in the graph, and try to merge vertex pair of largest penetration depth to make them even loops all_loops = [] visited = set() def dfs(node, parent, start_node, path): visited.add(node) path.append(node) for neighbor in graph[node]: if neighbor == parent: continue if neighbor == start_node and len(path) >= 3: all_loops.append(path.copy()) elif neighbor not in visited: dfs(neighbor, node, start_node, path) path.pop() visited.remove(node) for node in graph: dfs(node, None, node, []) visited.add(node) unique_loops = [] loop_keys = set() for loop in all_loops: loop_key = frozenset(loop) if len(loop) > 2 and loop_key not in loop_keys: loop_keys.add(loop_key) unique_loops.append(loop) # print(f'[INFO] find loop: {loop}') # convert loops to list of edges loop_edges = [] for loop in unique_loops: edges = set() for i in range(len(loop)): j = (i + 1) % len(loop) edges.add(tuple(sorted([loop[i], loop[j]]))) loop_edges.append(edges) # merge vertex pair of largest penetration depth to make them even loops ds = NamedDisjointSet(list(meshes.keys())) while True: # have to loop multiple times to make sure there is no odd loop... for i in range(len(loop_edges)): edges = loop_edges[i] if len(edges) % 2 == 0: # even continue # find the largest penetration depth max_penetration_depth = 0 max_edge = None for edge in edges: if penetration_depths[edge] > max_penetration_depth: max_penetration_depth = penetration_depths[edge] max_edge = edge # remove this edge from other loops if it exists (this will affect other loops, even make already even loops become odd) for j in range(len(loop_edges)): if max_edge in loop_edges[j]: loop_edges[j].remove(max_edge) # merge the vertex pair of largest penetration depth # print(f'[INFO] merge {max_edge[0]} and {max_edge[1]}') ds.merge(max_edge[0], max_edge[1]) has_odd_loop = False for i in range(len(loop_edges)): edges = loop_edges[i] if len(edges) % 2 == 1: has_odd_loop = True break if not has_odd_loop: break # merge groups graph_new = graph.copy() groups = ds.get_groups() for group in groups.values(): if len(group) <= 1: continue # print(f'[INFO] merge group: {group}') new_name = "_".join(group) new_mesh = trimesh.util.concatenate(list(meshes[name] for name in group)) meshes[new_name] = new_mesh graph_new[new_name] = set() for name in group: del meshes[name] for neighbor in graph[name]: if neighbor not in group: graph_new[new_name].add(neighbor) graph_new[neighbor].remove(name) graph_new[neighbor].add(new_name) del graph_new[name] print(f"[INFO] after merging odd loops, num_meshes = {len(meshes)}") return meshes, graph_new def run(path): print(f"[INFO] processing {path}") mesh = trimesh.load(path) if not opt.force_cc and isinstance(mesh, trimesh.Scene) and len(mesh.geometry) > 1: print(f"[INFO] scene: {len(mesh.geometry)} meshes") scene = mesh else: if isinstance(mesh, trimesh.Scene): mesh = mesh.to_mesh() # use meshiki backend mesh = Mesh(mesh.vertices, mesh.faces, verbose=opt.verbose, clean=False) mesh.smart_group_components() scene = mesh.export_components_as_trimesh_scene() # use trimesh # meshes = mesh.split(only_watertight=False) # # print(meshes) # scene = trimesh.Scene() # for mesh in meshes: # scene.add_geometry(mesh) print(f"[INFO] mesh: {len(scene.geometry)} components") ### box normalize into [-1, 1] bounds = scene.bounds # [2, 3] center = scene.centroid # [3] # print(f'[INFO] center = {center}, bounds = {bounds}') scale = 0.95 * 1 / np.max(bounds[1] - bounds[0]) transform_normalize = np.eye(4) transform_normalize[:3, 3] = -center transform_normalize[:3, :3] = np.diag(np.array([scale, scale, scale])) # print(transform_normalize) scene.apply_transform(transform_normalize) ### apply transform to vertices meshes = {} scene_graph = scene.graph.to_flattened() for k, v in scene_graph.items(): name = v["geometry"] if name in scene.geometry and isinstance(scene.geometry[name], trimesh.Trimesh): transform = v["transform"] mesh: trimesh.Trimesh = scene.geometry[name].apply_transform(transform) # drop all textures since we only need geom for parted data mesh.visual = trimesh.visual.ColorVisuals() # clean up mesh.merge_vertices(merge_tex=True, merge_norm=True) mesh.update_faces(mesh.unique_faces() & mesh.nondegenerate_faces()) mesh.fix_normals() meshes[name] = mesh ### smart grouping to avoid too many single-layer surface or too small objects if not opt.no_smart_group: meshes = smart_grouping(meshes) ### stitch open boundaries to make each mesh watertight if not opt.no_stitch: for name, mesh in meshes.items(): stitch_nonwatertight_mesh(mesh) ### coloring # build an undirected collision graph manager = trimesh.collision.CollisionManager() for name, mesh in meshes.items(): # scale up the mesh a little bit to take count of the collision margin mesh_dilated = mesh.copy() if not opt.no_dilate: center = mesh_dilated.centroid vertices = mesh_dilated.vertices - center max_radius = np.max(np.linalg.norm(vertices, axis=-1)) scale = (max_radius + opt.dilate_size) / max_radius # print(f'[INFO] dilate {name} by {scale}') mesh_dilated.vertices = vertices * scale + center manager.add_object(name, mesh_dilated) is_collide, collide_pairs, collide_data = manager.in_collision_internal(return_names=True, return_data=True) graph = {name: set() for name in meshes.keys()} penetration_depths = {} for data in collide_data: name1, name2 = list(data.names) graph[name1].add(name2) graph[name2].add(name1) name_key = tuple(sorted([name1, name2])) penetration_depth = data.depth penetration_depths[name_key] = penetration_depth # merge odd loops if not opt.no_merge_odd_loops: # if the graph is too complex, we will skip since it takes forever num_edges = sum(len(edges) for edges in graph.values()) if num_edges > 100: print(f"[WARN] skip {path} because of too many edges: {num_edges}") else: meshes, graph = merge_odd_loops(meshes, graph, penetration_depths) if opt.verbose: print(graph) # sort objects by distance to center name_to_centers = {} for name, mesh in meshes.items(): vmin = np.min(mesh.vertices, axis=0) vmax = np.max(mesh.vertices, axis=0) name_to_centers[name] = (vmin + vmax) / 2 name_with_dist = [] # [(name, dist), ...] for name, mesh_center in name_to_centers.items(): dist = np.linalg.norm(mesh_center) name_with_dist.append((name, dist)) name_with_dist.sort(key=lambda x: x[1]) # we will start graph coloring from center to border name_to_color = {} queue = [] initial_color = 0 for name, dist in name_with_dist: if name not in name_to_color: name_to_color[name] = initial_color initial_color = 1 - initial_color queue.append(name) while len(queue) > 0: name = queue.pop(0) for neighbor in graph[name]: if neighbor not in name_to_color: name_to_color[neighbor] = 1 - name_to_color[name] if opt.verbose: print(f"[INFO] color {neighbor} with {name_to_color[neighbor]}") queue.append(neighbor) else: if name_to_color[neighbor] == name_to_color[name]: print(f"[WARN] {name} and {neighbor} have the same color!") # get the two parts mesh_color0 = [] mesh_color1 = [] for name, color in name_to_color.items(): if color == 0: mesh_color0.append(meshes[name]) else: mesh_color1.append(meshes[name]) ### convert to a single mesh and export as glb mesh_color0 = trimesh.util.concatenate(mesh_color0) mesh_color1 = trimesh.util.concatenate(mesh_color1) name = os.path.splitext(os.path.basename(path))[0] # export separately mesh_color0.export(f"{opt.workspace}/{name}_color0.obj") mesh_color1.export(f"{opt.workspace}/{name}_color1.obj") # export together (offsetted) mesh_color1.vertices += [0, 0, 1] mesh_all = trimesh.util.concatenate([mesh_color0, mesh_color1]) mesh_all.export(f"{opt.workspace}/{name}.obj") os.makedirs(opt.workspace, exist_ok=True) if os.path.isdir(opt.test_path): file_paths = glob.glob(os.path.join(opt.test_path, "*")) for path in tqdm.tqdm(file_paths): run(path) else: run(opt.test_path)