|
- import time
- import numpy as np
- import open3d as o3d
- import torch
- from tqdm.auto import tqdm
- from pytorch3d.ops.knn import knn_points
-
- from chamfer_distance import ChamferDistance
- from sklearn.neighbors import NearestNeighbors
- from sklearn.metrics import f1_score
-
-
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # input from TensorBase ?
-
-
- def eval_mesh(gt_mesh_path, pred_mesh_path, save_path, batch_size=128, thresh=1):
- pred_pts = np.array(o3d.io.read_point_cloud(pred_mesh_path).points) # 1
- gt_pts = np.array(o3d.io.read_point_cloud(gt_mesh_path).points) # 2
-
- pred_pts = torch.from_numpy(pred_pts).to(device)
- gt_pts = torch.from_numpy(gt_pts).to(device)
-
- num_points1 = pred_pts.size(0)
- num_points2 = gt_pts.size(0)
-
- num_batches1 = int(num_points1 / batch_size)
- num_batches2 = int(num_points2 / batch_size)
-
- chamfer_dist_sum1 = 0.
- chamfer_dist_sum2 = 0.
-
- tp, fp, fn = 0, 0, 0
-
- '''no batchify for another set'''
- for i in tqdm(range(num_batches1), desc='Calculating dist1: '):
- batch_start_pred = i * batch_size
- batch_end_pred = (i + 1) * batch_size
- batch_pred = pred_pts[batch_start_pred:batch_end_pred]
-
- # l2
- dist1 = torch.cdist(batch_pred, gt_pts)
- min_dist1, _ = torch.min(dist1, dim=1)
-
- # l1
-
- tp += torch.sum(min_dist1 <= thresh)
- fp += torch.sum(min_dist1 > thresh)
-
- chamfer_dist_sum1 += torch.sum(min_dist1)
-
- for j in tqdm(range(num_batches2), desc='Calculating dist2: '):
- batch_start_gt = j * batch_size
- batch_end_gt = (j + 1) * batch_size
- batch_gt = gt_pts[batch_start_gt:batch_end_gt]
-
- # l2
- dist2 = torch.cdist(batch_gt, pred_pts)
- min_dist2, _ = torch.min(dist2, dim=1)
- # l1
- # gt_nn = knn_points(batch_gt, pred_pts, lengths1=len(batch_gt), lengths2=num_points2, norm=1, K=1)
- # dist1 = gt_nn.dists[..., 0]
-
- chamfer_dist_sum2 += torch.sum(min_dist2)
-
- fn = num_points2 - tp
- # for i in range(num_batches):
- # # 从点云1和点云2中选择当前批次的点
- # batch_start_gt = i * batch_size
- # batch_end_gt = (i + 1) * batch_size
- # batch_gt = points1[batch_start_gt:batch_end_gt]
- #
- # batch_start2 = i * batch_size
- # batch_end2 = (i + 1) * batch_size
- # batch_points2 = points2[batch_start2:batch_end2]
- #
- # # 计算当前批次的Chamfer距离
- # dist1 = torch.cdist(batch_gt, points2)
- # dist2 = torch.cdist(batch_points2, points1)
- #
- # min_dist1, _ = torch.min(dist1, dim=1)
- # min_dist2, _ = torch.min(dist2, dim=1)
- #
- # chamfer_dist_sum += torch.sum(min_dist1) + torch.sum(min_dist2)
-
- chamfer_dist = (chamfer_dist_sum1 / num_points1) + (chamfer_dist_sum2 / num_points2)
- precision = tp / (tp + fp)
- recall = tp / (tp + fn)
- f_score = 2 * precision * recall / (precision + recall)
-
- # tp = np.sum(dist1 <= thresh)
- # fp = np.sum(dist1 > thresh)
- # fn = np.sum(dist1 > thresh)
-
- desc = f'chamfer_dist: {chamfer_dist.item()}, f_score: {f_score}'
- print(desc)
- with open(f'{save_path}', 'w') as f:
- f.write(desc)
-
-
- if __name__ == '__main__':
- eval_mesh('/root/autodl-tmp/model_db/SOL-TensoRF/log/Replica/mesh/mesh.ply',
- '/root/autodl-tmp/model_db/SOL-TensoRF/log/Replica/mesh/room0_sem_MLPFea_rawfeas_tinyvoxel_sdf_test6_sdf_orgin_size.ply',
- '/root/autodl-tmp/model_db/SOL-TensoRF/log/Replica/mesh/sdf_thresh1.txt')
-
- # '/root/autodl-tmp/model_db/SOL-TensoRF/log/Replica/mesh/room0_sem_MLPFea_rawfeas_tinyvoxel_nosdf2_alpha_orgin_size.ply',
- # '/root/autodl-tmp/model_db/SOL-TensoRF/log/Replica/mesh/no_sdf.txt')
|