13 changed files with 438 additions and 221 deletions
@ -1,112 +1,247 @@ |
|||||
import trimesh |
|
||||
import numpy as np |
import numpy as np |
||||
from brep2sdf.data.sampler import sample_zero_surface_points_and_normals |
|
||||
from brep2sdf.networks.network import gradient |
|
||||
import torch |
import torch |
||||
import logging |
import argparse |
||||
|
from skimage import measure |
||||
# 配置日志记录 |
import time |
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') |
import trimesh |
||||
|
import pickle |
||||
def position_loss(pred_sdfs: torch.Tensor) -> torch.Tensor: |
from trimesh.proximity import ProximityQuery, signed_distance, closest_point |
||||
"""位置损失函数""" |
from brep2sdf.utils.logger import logger |
||||
# 保持梯度流 |
|
||||
squared_diff = torch.pow(pred_sdfs, 2) |
|
||||
return torch.mean(squared_diff) |
|
||||
|
|
||||
def average_normal_error(normals1: torch.Tensor, normals2: torch.Tensor) -> torch.Tensor: |
def sample_from_obj(obj_file, num_samples): |
||||
""" |
""" |
||||
计算平均法向量误差 (NAE) |
从OBJ文件中采样点 |
||||
:param normals1: 形状为 (B, 3) 的法向量张量 |
:param obj_file: OBJ文件路径 |
||||
:param normals2: 形状为 (B, 3) 的法向量张量 |
:param num_samples: 采样点数量 |
||||
:return: NAE 值 |
:return: 采样点数组 (N, 3) |
||||
""" |
""" |
||||
dot_products = torch.sum(normals1 * normals2, dim=-1) |
mesh = trimesh.load(obj_file) |
||||
absolute_dot_products = torch.abs(dot_products) |
points, _ = mesh.sample(num_samples, return_index=True) |
||||
angle_errors = 1 - absolute_dot_products |
return points |
||||
return torch.mean(angle_errors) |
|
||||
|
|
||||
def |
def obj_to_sdf(obj_file, depth, box_size): |
||||
|
""" |
||||
|
将OBJ文件转换为SDF网格 |
||||
|
:param obj_file: OBJ文件路径 |
||||
|
:param depth: 网格深度 |
||||
|
:param box_size: 边界框大小 |
||||
|
:return: SDF值数组和网格坐标 |
||||
|
""" |
||||
|
# 创建网格 |
||||
|
points, xx, yy, zz = create_grid(depth, box_size) |
||||
|
print(1) |
||||
|
# 计算SDF |
||||
|
mesh = trimesh.load(obj_file) |
||||
|
print(2) |
||||
|
sdf = signed_distance(mesh, points) |
||||
|
print(3) |
||||
|
sdf_grid = sdf.reshape(xx.shape) |
||||
|
print(4) |
||||
|
return sdf_grid, xx, yy, zz |
||||
|
|
||||
|
def extract_surface(sdf, xx, yy, zz, method='MC', bbox_size=1.0, feature_angle=30.0, voxel_size=0.01): |
||||
|
""" |
||||
|
提取零表面 |
||||
|
:param sdf: SDF值三维数组 |
||||
|
:param xx/yy/zz: 网格坐标 |
||||
|
:param method: 提取方法（MC: Marching Cubes） |
||||
|
:return: 顶点和面片 |
||||
|
""" |
||||
|
if method == 'MC': |
||||
|
verts, faces, _, _ = measure.marching_cubes(sdf, level=0) |
||||
|
elif method == 'EMC': |
||||
|
from iso_algorithms import enhanced_marching_cubes |
||||
|
verts, faces = enhanced_marching_cubes( |
||||
|
sdf, |
||||
|
feature_angle=feature_angle, |
||||
|
gradient_direction='descent' |
||||
|
) |
||||
|
elif method == 'DC': |
||||
|
from iso_algorithms import dual_contouring |
||||
|
verts, faces = dual_contouring(sdf, voxel_size=voxel_size) |
||||
|
else: |
||||
|
raise ValueError(f"不支持的算法: {method}") |
||||
|
|
||||
|
# 新增顶点后处理 |
||||
|
verts = (verts - sdf.shape[0]//2) / (sdf.shape[0]//2) / 2 * bbox_size # 归一化到[-1,1] |
||||
|
return verts, faces |
||||
|
|
||||
|
def save_ply(vertices, faces, filename): |
||||
|
""" |
||||
|
保存顶点和面片为PLY文件 |
||||
|
:param vertices: 顶点数组 (N, 3) |
||||
|
:param faces: 面片数组 (M, 3) |
||||
|
:param filename: 输出文件名 |
||||
|
""" |
||||
|
with open(filename, 'w') as f: |
||||
|
f.write("ply\n") |
||||
|
f.write("format ascii 1.0\n") |
||||
|
f.write(f"element vertex {len(vertices)}\n") |
||||
|
f.write("property float x\n") |
||||
|
f.write("property float y\n") |
||||
|
f.write("property float z\n") |
||||
|
f.write(f"element face {len(faces)}\n") |
||||
|
f.write("property list uchar int vertex_indices\n") |
||||
|
f.write("end_header\n") |
||||
|
for v in vertices: |
||||
|
f.write(f"{v[0]} {v[1]} {v[2]}\n") |
||||
|
for face in faces: |
||||
|
f.write(f"3 {face[0]} {face[1]} {face[2]}\n") |
||||
|
|
||||
|
|
||||
|
|
||||
|
def compute_iou(sdf1, sdf2, threshold=0.0): |
||||
|
""" |
||||
|
计算两个 SDF 之间的 IoU |
||||
|
:param sdf1: 第一个 SDF 数组 |
||||
|
:param sdf2: 第二个 SDF 数组 |
||||
|
:param threshold: 阈值，用于判断点是否在表面内 |
||||
|
:return: IoU 值 |
||||
|
""" |
||||
|
if sdf1.shape != sdf2.shape: |
||||
|
raise ValueError("sdf1 和 sdf2 的形状必须一致") |
||||
|
|
||||
|
inside1 = sdf1 <= threshold |
||||
|
inside2 = sdf2 <= threshold |
||||
|
intersection = np.logical_and(inside1, inside2).sum() |
||||
|
union = np.logical_or(inside1, inside2).sum() |
||||
|
|
||||
|
if union == 0: |
||||
|
logger.warning("union 为 0，无法计算 IoU") |
||||
|
return 0.0 |
||||
|
|
||||
|
iou = intersection / union |
||||
|
return iou |
||||
|
|
||||
|
def create_grid(depth, box_size): |
||||
|
""" |
||||
|
创建规则的三维网格 |
||||
|
:param depth: 网格深度（分辨率） |
||||
|
:param box_size: 边界框大小 |
||||
|
:return: 网格点坐标和网格坐标数组 |
||||
|
""" |
||||
|
grid_size = 2**depth + 1 |
||||
|
x = np.linspace(-box_size/2, box_size/2, grid_size) |
||||
|
y = np.linspace(-box_size/2, box_size/2, grid_size) |
||||
|
z = np.linspace(-box_size/2, box_size/2, grid_size) |
||||
|
xx, yy, zz = np.meshgrid(x, y, z, indexing='ij') |
||||
|
points = np.stack([xx.ravel(), yy.ravel(), zz.ravel()], axis=1) |
||||
|
return points, xx, yy, zz |
||||
|
|
||||
|
def load_brep_file(brep_path): |
||||
|
with open(brep_path, 'rb') as f: |
||||
|
brep_raw = pickle.load(f) |
||||
|
return brep_raw |
||||
|
|
||||
|
|
||||
|
|
||||
|
def test(): |
||||
|
data = load_brep_file("/home/wch/brep2sdf/data/output_data/00000003.xyz") |
||||
|
surfs = data["sampled_points_normals_sdf"] # 获取采样点数据 |
||||
|
points = surfs[:, :3] # 提取点坐标 |
||||
|
normals = surfs[:, 3:6] # 提取法向量 |
||||
|
sdf_values = surfs[:, 6] # 提取SDF值 |
||||
|
|
||||
|
print(len(points)) |
||||
|
print(f"points max: {np.max(points)}") # 打印points的最大值 |
||||
|
print(f"sdf_values max: {np.max(sdf_values)}") # 打印sdf_values的最大值 |
||||
|
|
||||
|
# 将点坐标和SDF值转换为网格格式 |
||||
|
grid_size = int(np.cbrt(len(points))) # 假设采样点是立方体网格 |
||||
|
sdf_grid = sdf_values.reshape((grid_size, grid_size, grid_size)) |
||||
|
|
||||
|
# 使用Marching Cubes提取零表面 |
||||
|
verts, faces, _, _ = measure.marching_cubes(sdf_grid, level=0) |
||||
|
# 打印verts和faces的最大值 |
||||
|
print(f"verts max: {np.max(verts)}") |
||||
|
print(f"faces max: {np.max(faces)}") |
||||
|
# 保存提取的网格 |
||||
|
save_ply(verts, faces, "output_mc.ply") |
||||
|
print("Marching Cubes 提取的网格已保存为 output_mc.ply") |
||||
|
|
||||
|
def sample_grid(trimesh_mesh_ncs: trimesh.Trimesh): |
||||
|
""" |
||||
|
np.ndarray | None: 形状为 (N, 7) 的数组 [x, y, z, nx, ny, nz, sdf]， |
||||
|
如果采样或计算失败则返回 None。 |
||||
|
""" |
||||
|
grid_size = 2**5 + 1 |
||||
|
start = -1 |
||||
|
end = 1 |
||||
|
x = np.linspace(start, end, grid_size) |
||||
|
y = np.linspace(start, end, grid_size) |
||||
|
z = np.linspace(start, end, grid_size) |
||||
|
xx, yy, zz = np.meshgrid(x, y, z, indexing='ij') |
||||
|
points = np.stack([xx.ravel(), yy.ravel(), zz.ravel()], axis=1) |
||||
|
|
||||
# ========== |
|
||||
def load_model(model_path): |
|
||||
"""加载模型的通用函数""" |
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
||||
try: |
try: |
||||
model = torch.jit.load(model_path).to(device) |
# Step 1: Compute signed distance |
||||
logging.info(f"成功加载模型: {model_path}") |
sdf = signed_distance(trimesh_mesh_ncs, points) |
||||
return model |
print(sdf) |
||||
except Exception as e: |
|
||||
logging.error(f"加载模型 {model_path} 时出错: {e}") |
|
||||
return None |
|
||||
|
|
||||
|
# Step 2: Compute normals via closest face |
||||
|
_, _, face_indices = closest_point(trimesh_mesh_ncs, points) |
||||
|
normals = trimesh_mesh_ncs.face_normals[face_indices] |
||||
|
|
||||
#========== |
# Step 3: Concatenate into (N, 7) |
||||
def nh(model_path, points): |
data = np.hstack([points, normals, sdf[:, None]]) |
||||
model = load_model(model_path) |
return data |
||||
if model is None: |
|
||||
return None |
|
||||
try: |
|
||||
return model(points) |
|
||||
except Exception as e: |
|
||||
logging.error(f"调用 NH 模型时出错: {e}") |
|
||||
return None |
|
||||
|
|
||||
def mine(model_path, points): |
|
||||
model = load_model(model_path) |
|
||||
if model is None: |
|
||||
return None |
|
||||
try: |
|
||||
return model.forward_background(points) |
|
||||
except Exception as e: |
except Exception as e: |
||||
logging.error(f"调用 mine 模型时出错: {e}") |
print(f"Error computing SDF or normals: {e}") |
||||
return None |
return None |
||||
|
|
||||
def main(): |
def test2(obj_file): |
||||
# 替换为实际的 obj 文件路径 |
mesh = trimesh.load(obj_file) |
||||
obj_file_path = "/home/wch/brep2sdf/data/gt_mesh/00000031.obj" |
surfs = sample_grid(mesh) |
||||
model_path = "/home/wch/brep2sdf/data/output_data/00000031.pt" |
points = surfs[:, :3] # 提取点坐标 |
||||
nh_model = "/home/wch/NH-Rep/data/output_data/00000031_0_50k_model_h.pt" |
normals = surfs[:, 3:6] # 提取法向量 |
||||
|
sdf_values = surfs[:, 6] # 提取SDF值 |
||||
|
|
||||
try: |
# 将点坐标和SDF值转换为网格格式 |
||||
# 读取 obj 文件 |
grid_size = int(np.cbrt(len(points))) # 假设采样点是立方体网格 |
||||
mesh = trimesh.load_mesh(obj_file_path) |
sdf_grid = sdf_values.reshape((grid_size, grid_size, grid_size)) |
||||
logging.info(f"成功读取 OBJ 文件: {obj_file_path}") |
|
||||
except Exception as e: |
|
||||
logging.error(f"读取 OBJ 文件 {obj_file_path} 时出错: {e}") |
|
||||
return |
|
||||
|
|
||||
try: |
# 使用Marching Cubes提取零表面 |
||||
# 调用采样函数 |
verts, faces, _, _ = measure.marching_cubes(sdf_grid, level=0) |
||||
result1 = sample_zero_surface_points_and_normals(mesh, num_samples=4096) |
|
||||
if result1 is None: |
# 打印verts和faces的最大值 |
||||
logging.error("采样失败，返回 None") |
print(f"verts max: {np.max(verts)}") |
||||
return |
print(f"faces max: {np.max(faces)}") |
||||
# 提取前 3 列作为坐标点 |
|
||||
coordinates = result1[:, :3] |
|
||||
# 将 ndarray 转换为 Tensor 并移动到设备上 |
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
||||
coordinates_tensor = torch.from_numpy(coordinates).float().to(device) |
|
||||
|
|
||||
sdf1 = nh(nh_model, coordinates_tensor) / 2 |
|
||||
sdf2 = mine(model_path, coordinates_tensor) |
|
||||
|
|
||||
if sdf1 is not None and sdf2 is not None: |
|
||||
loss1_ = position_loss(sdf1) |
|
||||
loss2 = position_loss(sdf2) |
|
||||
logging.info(f"NH 模型位置损失: {loss1.item()}") |
|
||||
logging.info(f"Mine 模型位置损失: {loss2.item()}") |
|
||||
|
|
||||
gt_normal = result1[:, 3:6] |
|
||||
normal1 = gradient(coordinates, sdf1) |
|
||||
normal2 = gradient(coordinates, sdf2) |
|
||||
nae1=average_normal_error(gt_normal, normal1) |
|
||||
nae2=average_normal_error(gt_normal, normal2) |
|
||||
else: |
|
||||
logging.error("无法计算损失，SDF 结果为 None") |
|
||||
|
|
||||
except Exception as e: |
# 保存提取的网格 |
||||
logging.error(f"处理过程中出现错误: {e}") |
save_ply(verts, faces, "output.ply") |
||||
|
print("Marching Cubes 提取的网格已保存为 output.ply") |
||||
|
|
||||
|
def main(): |
||||
|
parser = argparse.ArgumentParser(description='IsoSurface Generator') |
||||
|
parser.add_argument('-i', '--input', type=str, required=True, help='Input OBJ file') |
||||
|
parser.add_argument('-o', '--output', type=str, required=True, help='Output mesh file (.ply)') |
||||
|
parser.add_argument('--depth', type=int, default=5, help='网格深度（分辨率）') |
||||
|
parser.add_argument('--box_size', type=float, default=2.0, help='边界框大小') |
||||
|
parser.add_argument('--method', type=str, default='MC', choices=['MC', 'EMC', 'DC'], help='表面提取方法') |
||||
|
args = parser.parse_args() |
||||
|
|
||||
|
# 将OBJ转换为SDF网格 |
||||
|
sdf_grid, xx, yy, zz = obj_to_sdf(args.input, args.depth, args.box_size) |
||||
|
|
||||
|
# 提取表面 |
||||
|
print("Extracting surface...") |
||||
|
start_time = time.time() |
||||
|
verts, faces = extract_surface(sdf_grid, xx, yy, zz, args.method, args.box_size) |
||||
|
print(f"Surface extraction took {time.time() - start_time:.2f} seconds") |
||||
|
|
||||
|
# 保存网格 |
||||
|
save_ply(verts, faces, args.output) |
||||
|
print(f"Mesh saved to {args.output}") |
||||
|
|
||||
|
# 计算 IoU |
||||
|
sdf1 = sdf_grid # 第一个 SDF 网格 |
||||
|
sdf2 = sdf_grid # 第二个 SDF 网格（这里假设使用相同的网格进行计算） |
||||
|
iou = compute_iou(sdf1, sdf2) |
||||
|
print(f"IoU: {iou}") |
||||
|
|
||||
if __name__ == "__main__": |
if __name__ == "__main__": |
||||
main() |
#main() |
||||
|
#test() |
||||
|
test2("/home/wch/brep2sdf/data/gt_mesh/00000003.obj") |
||||
|
# python test.py -i /home/wch/brep2sdf/data/gt_mesh/00000003.obj -o output.ply --depth 6 --box_size 2.0 --method MC |
||||
Loading…
Reference in new issue