|
|
|
@ -1,5 +1,6 @@ |
|
|
|
import trimesh |
|
|
|
import numpy as np |
|
|
|
from brep2sdf.data.data import prepare_sdf_data,load_brep_file |
|
|
|
from brep2sdf.data.sampler import sample_zero_surface_points_and_normals |
|
|
|
from brep2sdf.utils.load import get_namelist, get_step_paths |
|
|
|
from brep2sdf.networks.network import gradient |
|
|
|
@ -7,6 +8,21 @@ import torch |
|
|
|
import os |
|
|
|
from brep2sdf.utils.logger import logger |
|
|
|
|
|
|
|
# 全局变量用于保存采样点 |
|
|
|
GLOBAL_SAMPLED_POINTS = None |
|
|
|
|
|
|
|
def sample_grid_points(mesh, nh_mesh, our_mesh, num_samples_iou): |
|
|
|
global GLOBAL_SAMPLED_POINTS |
|
|
|
if GLOBAL_SAMPLED_POINTS is not None: |
|
|
|
return GLOBAL_SAMPLED_POINTS |
|
|
|
# 从一个较大的空间范围采样点 |
|
|
|
bounds = np.vstack([mesh.bounds, nh_mesh.bounds, our_mesh.bounds]) |
|
|
|
min_bound = np.min(bounds, axis=0) |
|
|
|
max_bound = np.max(bounds, axis=0) |
|
|
|
points = np.random.uniform(min_bound, max_bound, (num_samples_iou, 3)) |
|
|
|
GLOBAL_SAMPLED_POINTS = points |
|
|
|
return points |
|
|
|
|
|
|
|
def position_loss(pred_sdfs: torch.Tensor) -> torch.Tensor: |
|
|
|
"""位置损失函数""" |
|
|
|
# 保持梯度流 |
|
|
|
@ -25,8 +41,54 @@ def average_normal_error(normals1: torch.Tensor, normals2: torch.Tensor) -> torc |
|
|
|
angle_errors = 1 - absolute_dot_products |
|
|
|
return torch.mean(angle_errors) |
|
|
|
|
|
|
|
def chamfer_distance(points_A: torch.Tensor, points_B: torch.Tensor) -> torch.Tensor: |
|
|
|
""" |
|
|
|
计算两个点集之间的 Chamfer Distance (CD) |
|
|
|
:param points_A: 形状为 (N, 3) 的点集 A |
|
|
|
:param points_B: 形状为 (M, 3) 的点集 B |
|
|
|
:return: CD 值 |
|
|
|
""" |
|
|
|
N = points_A.shape[0] |
|
|
|
M = points_B.shape[0] |
|
|
|
points_A_expanded = points_A.unsqueeze(1).expand(N, M, 3) |
|
|
|
points_B_expanded = points_B.unsqueeze(0).expand(N, M, 3) |
|
|
|
distances = torch.sum((points_A_expanded - points_B_expanded) ** 2, dim=-1) # (N, M) |
|
|
|
dist_A_to_B = torch.min(distances, dim=1)[0] # (N,) |
|
|
|
dist_B_to_A = torch.min(distances, dim=0)[0] # (M,) |
|
|
|
return (torch.mean(dist_A_to_B) + torch.mean(dist_B_to_A)) / 2 |
|
|
|
|
|
|
|
def hausdorff_distance(points_A: torch.Tensor, points_B: torch.Tensor) -> torch.Tensor: |
|
|
|
""" |
|
|
|
计算两个点集之间的 Two-Side Hausdorff Distance (HD) |
|
|
|
:param points_A: 形状为 (N, 3) 的点集 A |
|
|
|
:param points_B: 形状为 (M, 3) 的点集 B |
|
|
|
:return: HD 值 |
|
|
|
""" |
|
|
|
N = points_A.shape[0] |
|
|
|
M = points_B.shape[0] |
|
|
|
points_A_expanded = points_A.unsqueeze(1).expand(N, M, 3) |
|
|
|
points_B_expanded = points_B.unsqueeze(0).expand(N, M, 3) |
|
|
|
distances = torch.sum((points_A_expanded - points_B_expanded) ** 2, dim=-1) # (N, M) |
|
|
|
dist_A_to_B = torch.min(distances, dim=1)[0] # (N,) |
|
|
|
dist_B_to_A = torch.min(distances, dim=0)[0] # (M,) |
|
|
|
return torch.max(torch.max(dist_A_to_B), torch.max(dist_B_to_A)) |
|
|
|
|
|
|
|
def compute_iou(sdf1, sdf2, threshold=0.0): |
|
|
|
""" |
|
|
|
计算两个 SDF 之间的 IoU |
|
|
|
:param sdf1: 第一个 SDF 数组 |
|
|
|
:param sdf2: 第二个 SDF 数组 |
|
|
|
:param threshold: 阈值，用于判断点是否在表面内 |
|
|
|
:return: IoU 值 |
|
|
|
""" |
|
|
|
inside1 = sdf1 <= threshold |
|
|
|
inside2 = sdf2 <= threshold |
|
|
|
intersection = np.logical_and(inside1, inside2).sum() |
|
|
|
union = np.logical_or(inside1, inside2).sum() |
|
|
|
iou = intersection / union if union > 0 else 0.0 |
|
|
|
return iou |
|
|
|
|
|
|
|
# load |
|
|
|
def load_model(model_path): |
|
|
|
"""加载模型的通用函数""" |
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
@ -38,6 +100,9 @@ def load_model(model_path): |
|
|
|
logger.error(f"加载模型 {model_path} 时出错: {e}") |
|
|
|
return None |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# model |
|
|
|
def nh(model_path, points): |
|
|
|
model = load_model(model_path) |
|
|
|
if model is None: |
|
|
|
@ -53,16 +118,22 @@ def mine(model_path, points): |
|
|
|
if model is None: |
|
|
|
return None |
|
|
|
try: |
|
|
|
return model.forward_background(points) |
|
|
|
return model(points) |
|
|
|
except Exception as e: |
|
|
|
logger.error(f"调用 mine 模型时出错: {e}") |
|
|
|
return None |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def run(name): |
|
|
|
# 替换为实际的 obj 文件路径 |
|
|
|
obj_file_path = f"/home/wch/brep2sdf/data/gt_mesh/{name}.obj" |
|
|
|
model_path = f"/home/wch/brep2sdf/data/output_data/{name}.pt" |
|
|
|
nh_model = f"/home/wch/NH-Rep/data_backup/output_data/extracted/output_data/{name}_0_50k_model_h.pt" |
|
|
|
ply_nh = f"/home/wch/NH-Rep/data_backup/output_data/extracted/output_data/{name}_0_50k_model_h_nh.ply" |
|
|
|
ply_our = f"/home/wch/brep2sdf/data/output_data/{name}.ply" |
|
|
|
npz_path = f"/home/wch/brep2sdf/data/output_data/{name}.xyz" |
|
|
|
num_samples=4096 |
|
|
|
|
|
|
|
# 检查文件是否存在 |
|
|
|
if not os.path.isfile(obj_file_path): |
|
|
|
@ -79,7 +150,7 @@ def run(name): |
|
|
|
|
|
|
|
try: |
|
|
|
# 调用采样函数 |
|
|
|
result1 = sample_zero_surface_points_and_normals(mesh, num_samples=4096) |
|
|
|
result1 = sample_zero_surface_points_and_normals(mesh, num_samples) |
|
|
|
if result1 is None: |
|
|
|
logger.error("采样失败，返回 None") |
|
|
|
return |
|
|
|
@ -98,29 +169,115 @@ def run(name): |
|
|
|
loss2["de"] = position_loss(sdf2).item() |
|
|
|
logger.info(f"NH 模型位置损失: {loss1}") |
|
|
|
logger.info(f"Mine 模型位置损失: {loss2}") |
|
|
|
|
|
|
|
|
|
|
|
# 将 gt_normal 转换为 torch.Tensor 并移动到设备上 |
|
|
|
gt_normal = torch.from_numpy(result1[:, 3:6]).float().to(device) |
|
|
|
# 假设 gradient 函数已正确导入 |
|
|
|
normal1 = gradient(coordinates_tensor, sdf1) |
|
|
|
normal2 = gradient(coordinates_tensor, sdf2) |
|
|
|
|
|
|
|
|
|
|
|
loss1["nae"] = average_normal_error(gt_normal, normal1).item() |
|
|
|
loss2["nae"] = average_normal_error(gt_normal, normal2).item() |
|
|
|
|
|
|
|
|
|
|
|
print("NH 模型的平均法向量误差 (NAE):", loss1["nae"]) |
|
|
|
print("Mine 模型的平均法向量误差 (NAE):", loss2["nae"]) |
|
|
|
|
|
|
|
return loss1, loss2 |
|
|
|
|
|
|
|
|
|
|
|
else: |
|
|
|
logger.error("无法计算损失，SDF 结果为 None") |
|
|
|
|
|
|
|
# 读取 ply 文件 |
|
|
|
try: |
|
|
|
nh_mesh = trimesh.load_mesh(ply_nh) |
|
|
|
our_mesh = trimesh.load_mesh(ply_our) |
|
|
|
logger.info(f"成功读取 PLY 文件: {ply_nh} 和 {ply_our}") |
|
|
|
except Exception as e: |
|
|
|
logger.error(f"读取 PLY 文件时出错: {e}") |
|
|
|
return |
|
|
|
|
|
|
|
# 从网格中采样点 |
|
|
|
nh_points = torch.from_numpy(nh_mesh.sample(num_samples)).float().to(device) |
|
|
|
our_points = torch.from_numpy(our_mesh.sample(num_samples)).float().to(device) |
|
|
|
|
|
|
|
# 确保 coordinates 是 torch.Tensor 类型 |
|
|
|
loss1["cd"] = chamfer_distance(coordinates_tensor, nh_points).item() |
|
|
|
loss2["cd"] = chamfer_distance(coordinates_tensor, our_points).item() |
|
|
|
loss1["hd"] = hausdorff_distance(coordinates_tensor, nh_points).item() |
|
|
|
loss2["hd"] = hausdorff_distance(coordinates_tensor, our_points).item() |
|
|
|
|
|
|
|
# fea |
|
|
|
data = load_brep_file(npz_path) |
|
|
|
sampled_pnts=prepare_sdf_data(data["surf_ncs"],normals=data["surf_pnt_normals"],max_points=num_samples) |
|
|
|
|
|
|
|
# 展平处理 |
|
|
|
flattened_pnts = sampled_pnts.flatten() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# 修改此处，使用 clone().detach() |
|
|
|
if isinstance(flattened_pnts[0:3], torch.Tensor): |
|
|
|
f_pnts = flattened_pnts[0:3].clone().detach().to(device).view(-1, 3) |
|
|
|
else: |
|
|
|
f_pnts = torch.from_numpy(flattened_pnts[0:3]).clone().detach().to(device).view(-1, 3) |
|
|
|
|
|
|
|
if isinstance(flattened_pnts[3:6], torch.Tensor): |
|
|
|
f_normals = flattened_pnts[3:6].clone().detach().to(device).view(-1, 3) |
|
|
|
else: |
|
|
|
f_normals = torch.from_numpy(flattened_pnts[3:6]).clone().detach().to(device).view(-1, 3) |
|
|
|
|
|
|
|
# 检查 f_pnts 和 f_normals 的形状 |
|
|
|
if f_pnts.shape[-1] != 3 or f_normals.shape[-1] != 3: |
|
|
|
logger.error(f"f_pnts 形状: {f_pnts.shape}, f_normals 形状: {f_normals.shape}，期望最后一维尺寸为 3") |
|
|
|
return |
|
|
|
loss1["fcd"] = chamfer_distance(f_pnts, nh_points).item() |
|
|
|
loss2["fcd"] = chamfer_distance(f_pnts, our_points).item() |
|
|
|
loss1["fae"] = hausdorff_distance(f_normals, nh_points).item() |
|
|
|
loss2["fae"] = hausdorff_distance(f_normals, our_points).item() |
|
|
|
|
|
|
|
# 计算 IoU，从obj文件计算 |
|
|
|
# ... existing code ... |
|
|
|
|
|
|
|
# 计算 IoU，使用采样点方法 |
|
|
|
try: |
|
|
|
num_samples_iou = 10000 # 采样点数量，可以根据需要调整 |
|
|
|
# 调用封装的采样函数 |
|
|
|
points = sample_grid_points(mesh, nh_mesh, our_mesh, num_samples_iou) |
|
|
|
|
|
|
|
# 判断点是否在各个网格内部 |
|
|
|
inside_mesh = mesh.contains(points) |
|
|
|
inside_nh = nh_mesh.contains(points) |
|
|
|
inside_our = our_mesh.contains(points) |
|
|
|
|
|
|
|
# 计算 nh_mesh 与 mesh 的交集和并集 |
|
|
|
intersection_nh = np.logical_and(inside_mesh, inside_nh).sum() |
|
|
|
union_nh = np.logical_or(inside_mesh, inside_nh).sum() |
|
|
|
|
|
|
|
# 计算 our_mesh 与 mesh 的交集和并集 |
|
|
|
intersection_our = np.logical_and(inside_mesh, inside_our).sum() |
|
|
|
union_our = np.logical_or(inside_mesh, inside_our).sum() |
|
|
|
|
|
|
|
# 计算 IoU |
|
|
|
iou_nh = intersection_nh / union_nh if union_nh > 0 else 0.0 |
|
|
|
iou_our = intersection_our / union_our if union_our > 0 else 0.0 |
|
|
|
|
|
|
|
loss1["iou"] = iou_nh |
|
|
|
loss2["iou"] = iou_our |
|
|
|
except Exception as e: |
|
|
|
print(f"使用采样点计算 IoU 时出错: {e}") |
|
|
|
|
|
|
|
|
|
|
|
return loss1, loss2 |
|
|
|
|
|
|
|
except Exception as e: |
|
|
|
logger.error(f"处理过程中出现错误: {e}") |
|
|
|
|
|
|
|
def main(): |
|
|
|
names = get_namelist("/home/wch/brep2sdf/data/name_list.txt") |
|
|
|
tl1_de, tl1_nae, tl2_de, tl2_nae = 0.0, 0.0, 0.0, 0.0 |
|
|
|
tl1_cd, tl1_hd, tl2_cd, tl2_hd = 0.0, 0.0, 0.0, 0.0 |
|
|
|
tl1_fcd, tl1_fae, tl2_fcd, tl2_fae = 0.0, 0.0, 0.0, 0.0 |
|
|
|
# 新增累加 IoU 的变量 |
|
|
|
tl1_iou, tl2_iou = 0.0, 0.0 |
|
|
|
valid_count = 0 |
|
|
|
for name in names: |
|
|
|
result = run(name) |
|
|
|
@ -130,14 +287,61 @@ def main(): |
|
|
|
tl1_nae += l1["nae"] |
|
|
|
tl2_de += l2["de"] |
|
|
|
tl2_nae += l2["nae"] |
|
|
|
tl1_cd += l1["cd"] |
|
|
|
tl1_hd += l1["hd"] |
|
|
|
tl2_cd += l2["cd"] |
|
|
|
tl2_hd += l2["hd"] |
|
|
|
tl1_fcd += l1["fcd"] |
|
|
|
tl1_fae += l1["fae"] |
|
|
|
tl2_fcd += l2["fcd"] |
|
|
|
tl2_fae += l2["fae"] |
|
|
|
# 累加 IoU 的值 |
|
|
|
tl1_iou += l1["iou"] |
|
|
|
tl2_iou += l2["iou"] |
|
|
|
valid_count += 1 |
|
|
|
if valid_count > 0: |
|
|
|
print(f"NH 模型平均位置损失 (de): {tl1_de/valid_count}") |
|
|
|
print(f"NH 模型平均法向量误差 (nae): {tl1_nae/valid_count}") |
|
|
|
print(f"Mine 模型平均位置损失 (de): {tl2_de/valid_count}") |
|
|
|
print(f"Mine 模型平均法向量误差 (nae): {tl2_nae/valid_count}") |
|
|
|
avg_l1_de = tl1_de / valid_count |
|
|
|
avg_l1_nae = tl1_nae / valid_count |
|
|
|
avg_l2_de = tl2_de / valid_count |
|
|
|
avg_l2_nae = tl2_nae / valid_count |
|
|
|
avg_l1_cd = tl1_cd / valid_count |
|
|
|
avg_l1_hd = tl1_hd / valid_count |
|
|
|
avg_l2_cd = tl2_cd / valid_count |
|
|
|
avg_l2_hd = tl2_hd / valid_count |
|
|
|
avg_l1_fcd = tl1_fcd / valid_count |
|
|
|
avg_l1_fae = tl1_fae / valid_count |
|
|
|
avg_l2_fcd = tl2_fcd / valid_count |
|
|
|
avg_l2_fae = tl2_fae / valid_count |
|
|
|
# 计算 IoU 的平均值 |
|
|
|
avg_l1_iou = tl1_iou / valid_count |
|
|
|
avg_l2_iou = tl2_iou / valid_count |
|
|
|
# 打印表格表头 |
|
|
|
print("| 模型 | Chamfer Distance (CD) | Hausdorff Distance (HD) | 平均法向量误差 (NAE) | Feature Chamfer Distance (FCD) | Feature Angle Error (FAE) | 位置损失 (DE) | IoU |") |
|
|
|
print("|------|-----------------------|-------------------------|----------------------|--------------------------------|---------------------------|--------------|-----|") |
|
|
|
# 打印 NH 模型数据 |
|
|
|
print(f"| NH 模型 | {avg_l1_cd} | {avg_l1_hd} | {avg_l1_nae} | {avg_l1_fcd} | {avg_l1_fae} | {avg_l1_de} | {avg_l1_iou} |") |
|
|
|
# 打印 Mine 模型数据 |
|
|
|
print(f"| Mine 模型 | {avg_l2_cd} | {avg_l2_hd} | {avg_l2_nae} | {avg_l2_fcd} | {avg_l2_fae} | {avg_l2_de} | {avg_l2_iou} |") |
|
|
|
else: |
|
|
|
print("没有有效的结果，无法计算平均值。") |
|
|
|
|
|
|
|
def test(name_id): |
|
|
|
result = run(name_id) # 修正参数使用错误 |
|
|
|
if result is not None: |
|
|
|
l1, l2 = result |
|
|
|
# 打印表格表头 |
|
|
|
print("| 模型 | Chamfer Distance (CD) | Hausdorff Distance (HD) | 平均法向量误差 (NAE) | Feature Chamfer Distance (FCD) | Feature Angle Error (FAE) | 位置损失 (DE) | IoU |") |
|
|
|
print("|------|-----------------------|-------------------------|----------------------|--------------------------------|---------------------------|--------------|-----|") |
|
|
|
# 假设 IoU 数据存在于 l1 和 l2 中，如果不存在可以先忽略或者设置默认值 |
|
|
|
# 打印 NH 模型数据 |
|
|
|
print(f"| NH 模型 | {l1['cd']} | {l1['hd']} | {l1['nae']} | {l1['fcd']} | {l1['fae']} | {l1['de']} | {l1['iou']} |") |
|
|
|
# 打印 Mine 模型数据 |
|
|
|
print(f"| Mine 模型 | {l2['cd']} | {l2['hd']} | {l2['nae']} | {l2['fcd']} | {l2['fae']} | {l2['de']} | {l2['iou']} |") |
|
|
|
else: |
|
|
|
print("没有有效的结果。") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
main() |
|
|
|
#main() |
|
|
|
test("00000031") |
