import os import sys import torch import torch.nn as nn import numpy as np import argparse import trimesh from sklearn.decomposition import PCA import fpsample from tqdm import tqdm import threading import random # from tqdm.notebook import tqdm import time import copy import shutil from pathlib import Path from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed from collections import defaultdict import numba from numba import njit sys.path.append("..") from model import build_P3SAM, load_state_dict from utils.chamfer3D.dist_chamfer_3D import chamfer_3DDist cmd_loss = chamfer_3DDist() class P3SAM(nn.Module): def __init__(self): super().__init__() build_P3SAM(self) def load_state_dict(self, ckpt_path=None, state_dict=None, strict=True, assign=False, ignore_seg_mlp=False, ignore_seg_s2_mlp=False, ignore_iou_mlp=False): load_state_dict(self, ckpt_path=ckpt_path, state_dict=state_dict, strict=strict, assign=assign, ignore_seg_mlp=ignore_seg_mlp, ignore_seg_s2_mlp=ignore_seg_s2_mlp, ignore_iou_mlp=ignore_iou_mlp) def forward(self, feats, points, point_prompt, iter=1): """ feats: [K, N, 512] points: [K, N, 3] point_prompt: [K, N, 3] """ # print(feats.shape, points.shape, point_prompt.shape) point_num = points.shape[1] feats = feats.transpose(0, 1) # [N, K, 512] points = points.transpose(0, 1) # [N, K, 3] point_prompt = point_prompt.transpose(0, 1) # [N, K, 3] feats_seg = torch.cat([feats, points, point_prompt], dim=-1) # [N, K, 512+3+3] # 预测mask stage-1 pred_mask_1 = self.seg_mlp_1(feats_seg).squeeze(-1) # [N, K] pred_mask_2 = self.seg_mlp_2(feats_seg).squeeze(-1) # [N, K] pred_mask_3 = self.seg_mlp_3(feats_seg).squeeze(-1) # [N, K] pred_mask = torch.stack( [pred_mask_1, pred_mask_2, pred_mask_3], dim=-1 ) # [N, K, 3] for _ in range(iter): # 预测mask stage-2 feats_seg_2 = torch.cat([feats_seg, pred_mask], dim=-1) # [N, K, 512+3+3+3] feats_seg_global = self.seg_s2_mlp_g(feats_seg_2) # [N, K, 512] feats_seg_global = torch.max(feats_seg_global, dim=0).values # [K, 512] feats_seg_global = feats_seg_global.unsqueeze(0).repeat( point_num, 1, 1 ) # [N, K, 512] feats_seg_3 = torch.cat( [feats_seg_global, feats_seg_2], dim=-1 ) # [N, K, 512+3+3+3+512] pred_mask_s2_1 = self.seg_s2_mlp_1(feats_seg_3).squeeze(-1) # [N, K] pred_mask_s2_2 = self.seg_s2_mlp_2(feats_seg_3).squeeze(-1) # [N, K] pred_mask_s2_3 = self.seg_s2_mlp_3(feats_seg_3).squeeze(-1) # [N, K] pred_mask_s2 = torch.stack( [pred_mask_s2_1, pred_mask_s2_2, pred_mask_s2_3], dim=-1 ) # [N,, K 3] pred_mask = pred_mask_s2 mask_1 = torch.sigmoid(pred_mask_s2_1).to(dtype=torch.float32) # [N, K] mask_2 = torch.sigmoid(pred_mask_s2_2).to(dtype=torch.float32) # [N, K] mask_3 = torch.sigmoid(pred_mask_s2_3).to(dtype=torch.float32) # [N, K] feats_iou = torch.cat( [feats_seg_global, feats_seg, pred_mask_s2], dim=-1 ) # [N, K, 512+3+3+3+512] feats_iou = self.iou_mlp(feats_iou) # [N, K, 512] feats_iou = torch.max(feats_iou, dim=0).values # [K, 512] pred_iou = self.iou_mlp_out(feats_iou) # [K, 3] pred_iou = torch.sigmoid(pred_iou).to(dtype=torch.float32) # [K, 3] mask_1 = mask_1.transpose(0, 1) # [K, N] mask_2 = mask_2.transpose(0, 1) # [K, N] mask_3 = mask_3.transpose(0, 1) # [K, N] return mask_1, mask_2, mask_3, pred_iou def normalize_pc(pc): """ pc: (N, 3) """ max_, min_ = np.max(pc, axis=0), np.min(pc, axis=0) center = (max_ + min_) / 2 scale = (max_ - min_) / 2 scale = np.max(np.abs(scale)) pc = (pc - center) / (scale + 1e-10) return pc @torch.no_grad() def get_feat(model, points, normals): data_dict = { "coord": points, "normal": normals, "color": np.ones_like(points), "batch": np.zeros(points.shape[0], dtype=np.int64), } data_dict = model.transform(data_dict) for k in data_dict: if isinstance(data_dict[k], torch.Tensor): data_dict[k] = data_dict[k].cuda() point = model.sonata(data_dict) while "pooling_parent" in point.keys(): assert "pooling_inverse" in point.keys() parent = point.pop("pooling_parent") inverse = point.pop("pooling_inverse") parent.feat = torch.cat([parent.feat, point.feat[inverse]], dim=-1) point = parent feat = point.feat # [M, 1232] feat = model.mlp(feat) # [M, 512] feat = feat[point.inverse] # [N, 512] feats = feat return feats @torch.no_grad() def get_mask(model, feats, points, point_prompt, iter=1): """ feats: [N, 512] points: [N, 3] point_prompt: [K, 3] """ point_num = points.shape[0] prompt_num = point_prompt.shape[0] feats = feats.unsqueeze(1) # [N, 1, 512] feats = feats.repeat(1, prompt_num, 1).cuda() # [N, K, 512] points = torch.from_numpy(points).float().cuda().unsqueeze(1) # [N, 1, 3] points = points.repeat(1, prompt_num, 1) # [N, K, 3] prompt_coord = ( torch.from_numpy(point_prompt).float().cuda().unsqueeze(0) ) # [1, K, 3] prompt_coord = prompt_coord.repeat(point_num, 1, 1) # [N, K, 3] feats = feats.transpose(0, 1) # [K, N, 512] points = points.transpose(0, 1) # [K, N, 3] prompt_coord = prompt_coord.transpose(0, 1) # [K, N, 3] mask_1, mask_2, mask_3, pred_iou = model(feats, points, prompt_coord, iter) mask_1 = mask_1.transpose(0, 1) # [N, K] mask_2 = mask_2.transpose(0, 1) # [N, K] mask_3 = mask_3.transpose(0, 1) # [N, K] mask_1 = mask_1.detach().cpu().numpy() > 0.5 mask_2 = mask_2.detach().cpu().numpy() > 0.5 mask_3 = mask_3.detach().cpu().numpy() > 0.5 org_iou = pred_iou.detach().cpu().numpy() # [K, 3] return mask_1, mask_2, mask_3, org_iou def cal_iou(m1, m2): return np.sum(np.logical_and(m1, m2)) / np.sum(np.logical_or(m1, m2)) def cal_single_iou(m1, m2): return np.sum(np.logical_and(m1, m2)) / np.sum(m1) def iou_3d(box1, box2, signle=None): """ 计算两个三维边界框的交并比 (IoU) 参数: box1 (list): 第一个边界框的坐标 [x1_min, y1_min, z1_min, x1_max, y1_max, z1_max] box2 (list): 第二个边界框的坐标 [x2_min, y2_min, z2_min, x2_max, y2_max, z2_max] 返回: float: 交并比 (IoU) 值 """ # 计算交集的坐标 intersection_xmin = max(box1[0], box2[0]) intersection_ymin = max(box1[1], box2[1]) intersection_zmin = max(box1[2], box2[2]) intersection_xmax = min(box1[3], box2[3]) intersection_ymax = min(box1[4], box2[4]) intersection_zmax = min(box1[5], box2[5]) # 判断是否有交集 if ( intersection_xmin >= intersection_xmax or intersection_ymin >= intersection_ymax or intersection_zmin >= intersection_zmax ): return 0.0 # 无交集 # 计算交集的体积 intersection_volume = ( (intersection_xmax - intersection_xmin) * (intersection_ymax - intersection_ymin) * (intersection_zmax - intersection_zmin) ) # 计算两个盒子的体积 box1_volume = (box1[3] - box1[0]) * (box1[4] - box1[1]) * (box1[5] - box1[2]) box2_volume = (box2[3] - box2[0]) * (box2[4] - box2[1]) * (box2[5] - box2[2]) if signle is None: # 计算并集的体积 union_volume = box1_volume + box2_volume - intersection_volume elif signle == "1": union_volume = box1_volume elif signle == "2": union_volume = box2_volume else: raise ValueError("signle must be None or 1 or 2") # 计算 IoU iou = intersection_volume / union_volume if union_volume > 0 else 0.0 return iou def cal_point_bbox_iou(p1, p2, signle=None): min_p1 = np.min(p1, axis=0) max_p1 = np.max(p1, axis=0) min_p2 = np.min(p2, axis=0) max_p2 = np.max(p2, axis=0) box1 = [min_p1[0], min_p1[1], min_p1[2], max_p1[0], max_p1[1], max_p1[2]] box2 = [min_p2[0], min_p2[1], min_p2[2], max_p2[0], max_p2[1], max_p2[2]] return iou_3d(box1, box2, signle) def cal_bbox_iou(points, m1, m2): p1 = points[m1] p2 = points[m2] return cal_point_bbox_iou(p1, p2) def clean_mesh(mesh): """ mesh: trimesh.Trimesh """ # 1. 合并接近的顶点 mesh.merge_vertices() # 2. 删除重复的顶点 # 3. 删除重复的面片 mesh.process(True) return mesh def get_aabb_from_face_ids(mesh, face_ids): unique_ids = np.unique(face_ids) aabb = [] for i in unique_ids: if i == -1 or i == -2: continue _part_mask = face_ids == i _faces = mesh.faces[_part_mask] _faces = np.reshape(_faces, (-1)) _points = mesh.vertices[_faces] min_xyz = np.min(_points, axis=0) max_xyz = np.max(_points, axis=0) aabb.append([min_xyz, max_xyz]) return np.array(aabb) class Timer: def __init__(self, name): self.name = name def __enter__(self): self.start_time = time.time() return self # 可以返回 self 以便在 with 块内访问 def __exit__(self, exc_type, exc_val, exc_tb): self.end_time = time.time() self.elapsed_time = self.end_time - self.start_time print(f">>>>>>代码{self.name} 运行时间: {self.elapsed_time:.4f} 秒") def sample_points_pre_face(vertices, faces, n_point_per_face=2000): n_f = faces.shape[0] # 面片数量 # 生成随机数 u, v u = np.sqrt(np.random.rand(n_f, n_point_per_face, 1)) # (n_f, n_point_per_face, 1) v = np.random.rand(n_f, n_point_per_face, 1) # (n_f, n_point_per_face, 1) # 计算 barycentric 坐标 w0 = 1 - u w1 = u * (1 - v) w2 = u * v # (n_f, n_point_per_face, 1) # 从顶点中提取每个面的三个顶点 face_v_0 = vertices[faces[:, 0].reshape(-1)] # (n_f, 3) face_v_1 = vertices[faces[:, 1].reshape(-1)] # (n_f, 3) face_v_2 = vertices[faces[:, 2].reshape(-1)] # (n_f, 3) # 扩展维度以匹配 w0, w1, w2 的形状 face_v_0 = face_v_0.reshape(n_f, 1, 3) # (n_f, 1, 3) face_v_1 = face_v_1.reshape(n_f, 1, 3) # (n_f, 1, 3) face_v_2 = face_v_2.reshape(n_f, 1, 3) # (n_f, 1, 3) # 计算每个点的坐标 points = w0 * face_v_0 + w1 * face_v_1 + w2 * face_v_2 # (n_f, n_point_per_face, 3) return points def cal_cd_batch(p1, p2, pn=100000): p1_n = p1.shape[0] batch_num = (p1_n + pn - 1) // pn p2_cuda = torch.from_numpy(p2).cuda().float().unsqueeze(0) p1_cuda = torch.from_numpy(p1).cuda().float().unsqueeze(0) cd_res = [] for i in tqdm(range(batch_num)): start_idx = i * pn end_idx = min((i + 1) * pn, p1_n) _p1_cuda = p1_cuda[:, start_idx:end_idx, :] _, _, idx, _ = cmd_loss(_p1_cuda, p2_cuda) idx = idx[0].detach().cpu().numpy() cd_res.append(idx) cd_res = np.concatenate(cd_res, axis=0) return cd_res def remove_outliers_iqr(data, factor=1.5): """ 基于 IQR 去除离群值 :param data: 输入的列表或 NumPy 数组 :param factor: IQR 的倍数（默认 1.5） :return: 去除离群值后的列表 """ data = np.array(data, dtype=np.float32) q1 = np.percentile(data, 25) # 第一四分位数 q3 = np.percentile(data, 75) # 第三四分位数 iqr = q3 - q1 # 四分位距 lower_bound = q1 - factor * iqr upper_bound = q3 + factor * iqr return data[(data >= lower_bound) & (data <= upper_bound)].tolist() def better_aabb(points): x = points[:, 0] y = points[:, 1] z = points[:, 2] x = remove_outliers_iqr(x) y = remove_outliers_iqr(y) z = remove_outliers_iqr(z) min_xyz = np.array([np.min(x), np.min(y), np.min(z)]) max_xyz = np.array([np.max(x), np.max(y), np.max(z)]) return [min_xyz, max_xyz] def save_mesh(save_path, mesh, face_ids, color_map): face_colors = np.zeros((len(mesh.faces), 3), dtype=np.uint8) for i in tqdm(range(len(mesh.faces)), disable=True): _max_id = face_ids[i] if _max_id == -2: continue face_colors[i, :3] = color_map[_max_id] mesh_save = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces) mesh_save.visual.face_colors = face_colors mesh_save.export(save_path) mesh_save.export(save_path.replace(".glb", ".ply")) # print('保存mesh完成') scene_mesh = trimesh.Scene() scene_mesh.add_geometry(mesh_save) unique_ids = np.unique(face_ids) aabb = [] for i in unique_ids: if i == -1 or i == -2: continue _part_mask = face_ids == i _faces = mesh.faces[_part_mask] _faces = np.reshape(_faces, (-1)) _points = mesh.vertices[_faces] min_xyz, max_xyz = better_aabb(_points) center = (min_xyz + max_xyz) / 2 size = max_xyz - min_xyz box = trimesh.path.creation.box_outline() box.vertices *= size box.vertices += center box_color = np.array([[color_map[i][0], color_map[i][1], color_map[i][2], 255]]) box_color = np.repeat(box_color, len(box.entities), axis=0).astype(np.uint8) box.colors = box_color scene_mesh.add_geometry(box) min_xyz = np.min(_points, axis=0) max_xyz = np.max(_points, axis=0) aabb.append([min_xyz, max_xyz]) scene_mesh.export(save_path.replace(".glb", "_aabb.glb")) aabb = np.array(aabb) np.save(save_path.replace(".glb", "_aabb.npy"), aabb) np.save(save_path.replace(".glb", "_face_ids.npy"), face_ids) def mesh_sam( model, mesh, save_path, point_num=100000, prompt_num=400, save_mid_res=False, show_info=False, post_process=False, threshold=0.95, clean_mesh_flag=True, seed=42, prompt_bs=32, ): with Timer("加载mesh"): model, model_parallel = model if clean_mesh_flag: mesh = clean_mesh(mesh) mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, process=False) if show_info: print(f"点数：{mesh.vertices.shape[0]} 面片数：{mesh.faces.shape[0]}") point_num = 100000 prompt_num = 400 with Timer("采样点云"): _points, face_idx = trimesh.sample.sample_surface(mesh, point_num, seed=seed) _points_org = _points.copy() _points = normalize_pc(_points) normals = mesh.face_normals[face_idx] # _points = _points + np.random.normal(0, 1, size=_points.shape) * 0.01 # normals = normals * 0. # debug no normal if show_info: print(f"点数：{point_num} 面片数：{mesh.faces.shape[0]}") with Timer("获取特征"): _feats = get_feat(model, _points, normals) if show_info: print("预处理特征") if save_mid_res: feat_save = _feats.float().detach().cpu().numpy() data_scaled = feat_save / np.linalg.norm(feat_save, axis=-1, keepdims=True) pca = PCA(n_components=3) data_reduced = pca.fit_transform(data_scaled) data_reduced = (data_reduced - data_reduced.min()) / ( data_reduced.max() - data_reduced.min() ) _colors_pca = (data_reduced * 255).astype(np.uint8) pc_save = trimesh.points.PointCloud(_points, colors=_colors_pca) pc_save.export(os.path.join(save_path, "point_pca.glb")) pc_save.export(os.path.join(save_path, "point_pca.ply")) if show_info: print("PCA获取特征颜色") with Timer("FPS采样提示点"): fps_idx = fpsample.fps_sampling(_points, prompt_num) _point_prompts = _points[fps_idx] if save_mid_res: trimesh.points.PointCloud(_point_prompts, colors=_colors_pca[fps_idx]).export( os.path.join(save_path, "point_prompts_pca.glb") ) trimesh.points.PointCloud(_point_prompts, colors=_colors_pca[fps_idx]).export( os.path.join(save_path, "point_prompts_pca.ply") ) if show_info: print("采样完成") with Timer("推理"): bs = prompt_bs step_num = prompt_num // bs + 1 mask_res = [] iou_res = [] for i in tqdm(range(step_num), disable=not show_info): cur_propmt = _point_prompts[bs * i : bs * (i + 1)] pred_mask_1, pred_mask_2, pred_mask_3, pred_iou = get_mask( model_parallel, _feats, _points, cur_propmt ) pred_mask = np.stack( [pred_mask_1, pred_mask_2, pred_mask_3], axis=-1 ) # [N, K, 3] max_idx = np.argmax(pred_iou, axis=-1) # [K] for j in range(max_idx.shape[0]): mask_res.append(pred_mask[:, j, max_idx[j]]) iou_res.append(pred_iou[j, max_idx[j]]) mask_res = np.stack(mask_res, axis=-1) # [N, K] if show_info: print("prmopt 推理完成") with Timer("根据IOU排序"): iou_res = np.array(iou_res).tolist() mask_iou = [[mask_res[:, i], iou_res[i]] for i in range(prompt_num)] mask_iou_sorted = sorted(mask_iou, key=lambda x: x[1], reverse=True) mask_sorted = [mask_iou_sorted[i][0] for i in range(prompt_num)] iou_sorted = [mask_iou_sorted[i][1] for i in range(prompt_num)] # clusters = {} # for i in tqdm(range(prompt_num), desc="NMS", disable=not show_info): # _mask = mask_sorted[i] # union_flag = False # for j in clusters.keys(): # if cal_iou(_mask, mask_sorted[j]) > 0.9: # clusters[j].append(i) # union_flag = True # break # if not union_flag: # clusters[i] = [i] with Timer("NMS"): clusters = defaultdict(list) with ThreadPoolExecutor(max_workers=20) as executor: for i in tqdm(range(prompt_num), desc="NMS", disable=not show_info): _mask = mask_sorted[i] futures = [] for j in clusters.keys(): futures.append(executor.submit(cal_iou, _mask, mask_sorted[j])) for j, future in zip(clusters.keys(), futures): if future.result() > 0.9: clusters[j].append(i) break else: clusters[i].append(i) # print(clusters) if show_info: print(f"NMS完成，mask数量：{len(clusters)}") if save_mid_res: part_mask_save_path = os.path.join(save_path, "part_mask") if os.path.exists(part_mask_save_path): shutil.rmtree(part_mask_save_path) os.makedirs(part_mask_save_path, exist_ok=True) for i in tqdm(clusters.keys(), desc="保存mask", disable=not show_info): cluster_num = len(clusters[i]) cluster_iou = iou_sorted[i] cluster_area = np.sum(mask_sorted[i]) if cluster_num <= 2: continue mask_save = mask_sorted[i] mask_save = np.expand_dims(mask_save, axis=-1) mask_save = np.repeat(mask_save, 3, axis=-1) mask_save = (mask_save * 255).astype(np.uint8) point_save = trimesh.points.PointCloud(_points, colors=mask_save) point_save.export( os.path.join( part_mask_save_path, f"mask_{i}_iou_{cluster_iou:.5f}_area_{cluster_area:.5f}_num_{cluster_num}.glb", ) ) # 过滤只有一个mask的cluster with Timer("过滤只有一个mask的cluster"): filtered_clusters = [] other_clusters = [] for i in clusters.keys(): if len(clusters[i]) > 2: filtered_clusters.append(i) else: other_clusters.append(i) if show_info: print( f"过滤前：{len(clusters)} 个cluster，" f"过滤后：{len(filtered_clusters)} 个cluster" ) # 再次合并 with Timer("再次合并"): filtered_clusters_num = len(filtered_clusters) cluster2 = {} is_union = [False] * filtered_clusters_num for i in range(filtered_clusters_num): if is_union[i]: continue cur_cluster = filtered_clusters[i] cluster2[cur_cluster] = [cur_cluster] for j in range(i + 1, filtered_clusters_num): if is_union[j]: continue tar_cluster = filtered_clusters[j] # if cal_single_iou(mask_sorted[tar_cluster], mask_sorted[cur_cluster]) > 0.9: # if cal_iou(mask_sorted[tar_cluster], mask_sorted[cur_cluster]) > 0.5: if ( cal_bbox_iou( _points, mask_sorted[tar_cluster], mask_sorted[cur_cluster] ) > 0.5 ): cluster2[cur_cluster].append(tar_cluster) is_union[j] = True if show_info: print(f"再次合并，合并数量：{len(cluster2.keys())}") with Timer("计算没有mask的点"): no_mask = np.ones(point_num) for i in cluster2: part_mask = mask_sorted[i] no_mask[part_mask] = 0 if show_info: print( f"{np.sum(no_mask == 1)} 个点没有mask," f" 占比：{np.sum(no_mask == 1) / point_num:.4f}" ) with Timer("修补遗漏mask"): # 查询漏掉的mask for i in tqdm(range(len(mask_sorted)), desc="漏掉mask", disable=not show_info): if i in cluster2: continue part_mask = mask_sorted[i] _iou = cal_single_iou(part_mask, no_mask) if _iou > 0.7: cluster2[i] = [i] no_mask[part_mask] = 0 if save_mid_res: mask_save = mask_sorted[i] mask_save = np.expand_dims(mask_save, axis=-1) mask_save = np.repeat(mask_save, 3, axis=-1) mask_save = (mask_save * 255).astype(np.uint8) point_save = trimesh.points.PointCloud(_points, colors=mask_save) cluster_iou = iou_sorted[i] cluster_area = int(np.sum(mask_sorted[i])) cluster_num = 1 point_save.export( os.path.join( part_mask_save_path, f"mask_{i}_iou_{cluster_iou:.5f}_area_{cluster_area:.5f}_num_{cluster_num}.glb", ) ) # print(cluster2) # print(len(cluster2.keys())) if show_info: print(f"修补遗漏mask：{len(cluster2.keys())}") with Timer("计算点云最终mask"): final_mask = list(cluster2.keys()) final_mask_area = [int(np.sum(mask_sorted[i])) for i in final_mask] final_mask_area = [ [final_mask[i], final_mask_area[i]] for i in range(len(final_mask)) ] final_mask_area_sorted = sorted( final_mask_area, key=lambda x: x[1], reverse=True ) final_mask_sorted = [ final_mask_area_sorted[i][0] for i in range(len(final_mask_area)) ] final_mask_area_sorted = [ final_mask_area_sorted[i][1] for i in range(len(final_mask_area)) ] # print(final_mask_sorted) # print(final_mask_area_sorted) if show_info: print(f"最终mask数量：{len(final_mask_sorted)}") with Timer("点云上色"): # 生成color map color_map = {} for i in final_mask_sorted: part_color = np.random.rand(3) * 255 color_map[i] = part_color # print(color_map) result_mask = -np.ones(point_num, dtype=np.int64) for i in final_mask_sorted: part_mask = mask_sorted[i] result_mask[part_mask] = i if save_mid_res: # 保存点云结果 result_colors = np.zeros_like(_colors_pca) for i in final_mask_sorted: part_color = color_map[i] part_mask = mask_sorted[i] result_colors[part_mask, :3] = part_color trimesh.points.PointCloud(_points, colors=result_colors).export( os.path.join(save_path, "auto_mask_cluster.glb") ) trimesh.points.PointCloud(_points, colors=result_colors).export( os.path.join(save_path, "auto_mask_cluster.ply") ) if show_info: print("保存点云完成") with Timer("后处理"): valid_mask = result_mask >= 0 _org = _points_org[valid_mask] _results = result_mask[valid_mask] pre_face = 10 _face_points = sample_points_pre_face( mesh.vertices, mesh.faces, n_point_per_face=pre_face ) _face_points = np.reshape(_face_points, (len(mesh.faces) * pre_face, 3)) _idx = cal_cd_batch(_face_points, _org) _idx_res = _results[_idx] _idx_res = np.reshape(_idx_res, (-1, pre_face)) face_ids = [] for i in range(len(mesh.faces)): _label = np.argmax(np.bincount(_idx_res[i] + 2)) - 2 face_ids.append(_label) final_face_ids = np.array(face_ids) if save_mid_res: save_mesh( os.path.join(save_path, "auto_mask_mesh_final.glb"), mesh, final_face_ids, color_map, ) with Timer("计算最后的aabb"): aabb = get_aabb_from_face_ids(mesh, final_face_ids) return aabb, final_face_ids, mesh class AutoMask: def __init__( self, ckpt_path=None, point_num=100000, prompt_num=400, threshold=0.95, post_process=True, automask_instance=None, ): """ ckpt_path: str, 模型路径 point_num: int, 采样点数量 prompt_num: int, 提示数量 threshold: float, 阈值 post_process: bool, 是否后处理 """ if automask_instance is not None: self.model = automask_instance.model self.model_parallel = automask_instance.model_parallel else: self.model = P3SAM() self.model.load_state_dict(ckpt_path) self.model.eval() # self.model_parallel = torch.nn.DataParallel(self.model) self.model_parallel = self.model self.model.cuda() self.model_parallel.cuda() self.point_num = point_num self.prompt_num = prompt_num self.threshold = threshold self.post_process = post_process def predict_aabb( self, mesh, point_num=None, prompt_num=None, threshold=None, post_process=None, save_path=None, save_mid_res=False, show_info=True, clean_mesh_flag=True, seed=42, is_parallel=True, prompt_bs=32, ): """ Parameters: mesh: trimesh.Trimesh, 输入网格 point_num: int, 采样点数量 prompt_num: int, 提示数量 threshold: float, 阈值 post_process: bool, 是否后处理 Returns: aabb: np.ndarray, 包围盒 face_ids: np.ndarray, 面id """ point_num = point_num if point_num is not None else self.point_num prompt_num = prompt_num if prompt_num is not None else self.prompt_num threshold = threshold if threshold is not None else self.threshold post_process = post_process if post_process is not None else self.post_process return mesh_sam( [self.model, self.model_parallel if is_parallel else self.model], mesh, save_path=save_path, point_num=point_num, prompt_num=prompt_num, threshold=threshold, post_process=post_process, show_info=show_info, save_mid_res=save_mid_res, clean_mesh_flag=clean_mesh_flag, seed=seed, prompt_bs=prompt_bs, ) def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False if __name__ == "__main__": argparser = argparse.ArgumentParser() argparser.add_argument( "--ckpt_path", type=str, default=None, help="模型路径" ) argparser.add_argument( "--mesh_path", type=str, default="assets/1.glb", help="输入网格路径" ) argparser.add_argument( "--output_path", type=str, default="results/1", help="保存路径" ) argparser.add_argument("--point_num", type=int, default=100000, help="采样点数量") argparser.add_argument("--prompt_num", type=int, default=400, help="提示数量") argparser.add_argument("--threshold", type=float, default=0.95, help="阈值") argparser.add_argument("--post_process", type=int, default=0, help="是否后处理") argparser.add_argument( "--save_mid_res", type=int, default=1, help="是否保存中间结果" ) argparser.add_argument("--show_info", type=int, default=1, help="是否显示信息") argparser.add_argument( "--show_time_info", type=int, default=1, help="是否显示时间信息" ) argparser.add_argument("--seed", type=int, default=42, help="随机种子") argparser.add_argument("--parallel", type=int, default=1, help="是否使用多卡") argparser.add_argument( "--prompt_bs", type=int, default=32, help="提示点推理时的batch size大小" ) argparser.add_argument("--clean_mesh", type=int, default=1, help="是否清洗网格") args = argparser.parse_args() Timer.STATE = args.show_time_info output_path = args.output_path os.makedirs(output_path, exist_ok=True) ckpt_path = args.ckpt_path auto_mask = AutoMask(ckpt_path) mesh_path = args.mesh_path if os.path.isdir(mesh_path): for file in os.listdir(mesh_path): if not ( file.endswith(".glb") or file.endswith(".obj") or file.endswith(".ply") ): continue _mesh_path = os.path.join(mesh_path, file) _output_path = os.path.join(output_path, file[:-4]) os.makedirs(_output_path, exist_ok=True) mesh = trimesh.load(_mesh_path, force="mesh") set_seed(args.seed) aabb, face_ids, mesh = auto_mask.predict_aabb( mesh, save_path=_output_path, point_num=args.point_num, prompt_num=args.prompt_num, threshold=args.threshold, post_process=args.post_process, save_mid_res=args.save_mid_res, show_info=args.show_info, seed=args.seed, is_parallel=args.parallel, clean_mesh_flag=args.clean_mesh, ) else: mesh = trimesh.load(mesh_path, force="mesh") set_seed(args.seed) aabb, face_ids, mesh = auto_mask.predict_aabb( mesh, save_path=output_path, point_num=args.point_num, prompt_num=args.prompt_num, threshold=args.threshold, post_process=args.post_process, save_mid_res=args.save_mid_res, show_info=args.show_info, seed=args.seed, is_parallel=args.parallel, clean_mesh_flag=args.clean_mesh, ) ############################################### ## 可以通过以下代码保存返回的结果 ## You can save the returned result by the following code ################# save result ################# # color_map = {} # unique_ids = np.unique(face_ids) # for i in unique_ids: # if i == -1: # continue # part_color = np.random.rand(3) * 255 # color_map[i] = part_color # face_colors = [] # for i in face_ids: # if i == -1: # face_colors.append([0, 0, 0]) # else: # face_colors.append(color_map[i]) # face_colors = np.array(face_colors).astype(np.uint8) # mesh_save = mesh.copy() # mesh_save.visual.face_colors = face_colors # mesh_save.export(os.path.join(output_path, 'auto_mask_mesh.glb')) # scene_mesh = trimesh.Scene() # scene_mesh.add_geometry(mesh_save) # for i in range(len(aabb)): # min_xyz, max_xyz = aabb[i] # center = (min_xyz + max_xyz) / 2 # size = max_xyz - min_xyz # box = trimesh.path.creation.box_outline() # box.vertices *= size # box.vertices += center # scene_mesh.add_geometry(box) # scene_mesh.export(os.path.join(output_path, 'auto_mask_aabb.glb')) ################# save result ################# """ python auto_mask_no_postprocess.py --parallel 0 python auto_mask_no_postprocess.py --ckpt_path ../weights/p3sam.ckpt --mesh_path assets/1.glb --output_path results/1 --parallel 0 python auto_mask_no_postprocess.py --ckpt_path ../weights/p3sam.ckpt --mesh_path assets --output_path results/all_no_postprocess """