tohid4n/PartCrafter

1from src.utils.typing_utils import *2 3import trimesh4import numpy as np5from sklearn.neighbors import NearestNeighbors6 7def sample_from_mesh(8    mesh: trimesh.Trimesh,9    num_samples: Optional[int] = 10000,10):11    if num_samples is None:12        return mesh.vertices13    else:14        return mesh.sample(num_samples)15 16def sample_two_meshes(17    mesh1: trimesh.Trimesh,18    mesh2: trimesh.Trimesh,19    num_samples: Optional[int] = 10000,20):21    points1 = sample_from_mesh(mesh1, num_samples)22    points2 = sample_from_mesh(mesh2, num_samples)23    return points1, points224 25def compute_nearest_distance(26    points1: np.ndarray,27    points2: np.ndarray,28    metric: str = 'l2'29) -> np.ndarray:30    # Compute nearest neighbor distance from points1 to points231    nn = NearestNeighbors(n_neighbors=1, leaf_size=30, algorithm='kd_tree', metric=metric).fit(points2)32    min_dist = nn.kneighbors(points1)[0]33    return min_dist34 35def compute_mutual_nearest_distance(36    points1: np.ndarray,37    points2: np.ndarray,38    metric: str = 'l2'39) -> np.ndarray:40    min_1_to_2 = compute_nearest_distance(points1, points2, metric=metric)41    min_2_to_1 = compute_nearest_distance(points2, points1, metric=metric)42    return min_1_to_2, min_2_to_143 44def compute_mutual_nearest_distance_for_meshes(45    mesh1: trimesh.Trimesh,46    mesh2: trimesh.Trimesh,47    num_samples: Optional[int] = 10000,48    metric: str = 'l2'49) -> Tuple[np.ndarray, np.ndarray]:50    points1 = sample_from_mesh(mesh1, num_samples)51    points2 = sample_from_mesh(mesh2, num_samples)52    min_1_to_2, min_2_to_1 = compute_mutual_nearest_distance(points1, points2, metric=metric)53    return min_1_to_2, min_2_to_154 55def compute_chamfer_distance(56    mesh1: trimesh.Trimesh,57    mesh2: trimesh.Trimesh,58    num_samples: int = 10000,59    metric: str = 'l2'60):61    min_1_to_2, min_2_to_1 = compute_mutual_nearest_distance_for_meshes(mesh1, mesh2, num_samples, metric=metric)62    chamfer_dist = np.mean(min_2_to_1) + np.mean(min_1_to_2)63    return chamfer_dist64 65def compute_f_score(66    mesh1: trimesh.Trimesh,67    mesh2: trimesh.Trimesh,68    num_samples: int = 10000,69    threshold: float = 0.1,70    metric: str = 'l2'71):72    min_1_to_2, min_2_to_1 = compute_mutual_nearest_distance_for_meshes(mesh1, mesh2, num_samples, metric=metric)73    precision_1 = np.mean((min_1_to_2 < threshold).astype(np.float32))74    precision_2 = np.mean((min_2_to_1 < threshold).astype(np.float32))75    fscore = 2 * precision_1 * precision_2 / (precision_1 + precision_2)76    return fscore77 78def compute_cd_and_f_score(79    mesh1: trimesh.Trimesh,80    mesh2: trimesh.Trimesh,81    num_samples: Optional[int] = 10000,82    threshold: float = 0.1,83    metric: str = 'l2'84):85    min_1_to_2, min_2_to_1 = compute_mutual_nearest_distance_for_meshes(mesh1, mesh2, num_samples, metric=metric)86    chamfer_dist = np.mean(min_2_to_1) + np.mean(min_1_to_2)87    precision_1 = np.mean((min_1_to_2 < threshold).astype(np.float32))88    precision_2 = np.mean((min_2_to_1 < threshold).astype(np.float32))89    fscore = 2 * precision_1 * precision_2 / (precision_1 + precision_2)90    return chamfer_dist, fscore91 92def compute_cd_and_f_score_in_training(93    gt_surface: np.ndarray,94    pred_mesh: trimesh.Trimesh,95    num_samples: int = 204800,96    threshold: float = 0.1,97    metric: str = 'l2'98):99    gt_points = gt_surface[:, :3]100    num_samples = max(num_samples, gt_points.shape[0])101    gt_points = gt_points[np.random.choice(gt_points.shape[0], num_samples, replace=False)]102    pred_points = sample_from_mesh(pred_mesh, num_samples)103    min_1_to_2, min_2_to_1 = compute_mutual_nearest_distance(gt_points, pred_points, metric=metric)104    chamfer_dist = np.mean(min_2_to_1) + np.mean(min_1_to_2)105    precision_1 = np.mean((min_1_to_2 < threshold).astype(np.float32))106    precision_2 = np.mean((min_2_to_1 < threshold).astype(np.float32))107    fscore = 2 * precision_1 * precision_2 / (precision_1 + precision_2)108    return chamfer_dist, fscore109 110def get_voxel_set(111    mesh: trimesh.Trimesh,112    num_grids: int = 64,113    scale: float = 2.0,114):115    if not isinstance(mesh, trimesh.Trimesh):116        raise ValueError("mesh must be a trimesh.Trimesh object")117    pitch = scale / num_grids118    voxel_girds: trimesh.voxel.base.VoxelGrid = mesh.voxelized(pitch=pitch).fill()119    voxels = set(map(tuple, np.round(voxel_girds.points / pitch).astype(int)))120    return voxels121 122def compute_IoU(123    mesh1: trimesh.Trimesh,124    mesh2: trimesh.Trimesh,125    num_grids: int = 64,126    scale: float = 2.0,127):128    if not isinstance(mesh1, trimesh.Trimesh) or not isinstance(mesh2, trimesh.Trimesh):129        raise ValueError("mesh1 and mesh2 must be trimesh.Trimesh objects")130    voxels1 = get_voxel_set(mesh1, num_grids, scale)131    voxels2 = get_voxel_set(mesh2, num_grids, scale)132    intersection = voxels1 & voxels2133    union = voxels1 | voxels2134    iou = len(intersection) / len(union) if len(union) > 0 else 0.0135    return iou136 137def compute_IoU_for_scene(138    scene: Union[trimesh.Scene, List[trimesh.Trimesh]],139    num_grids: int = 64,140    scale: float = 2.0,141    return_type: Literal["iou", "iou_list"] = "iou",142):143    if isinstance(scene, trimesh.Scene):144        scene = scene.dump()145    if isinstance(scene, list) and len(scene) > 1 and isinstance(scene[0], trimesh.Trimesh):146        meshes = scene147    else:148        raise ValueError("scene must be a trimesh.Scene object or a list of trimesh.Trimesh objects")149    ious = []150    for i in range(len(meshes)):151        for j in range(i+1, len(meshes)):152            iou = compute_IoU(meshes[i], meshes[j], num_grids, scale)153            ious.append(iou)154    if return_type == "iou":155        return np.mean(ious)156    elif return_type == "iou_list":157        return ious158    else:159        raise ValueError("return_type must be 'iou' or 'iou_list'")