13 changed files with 438 additions and 221 deletions
@ -1,112 +1,247 @@ |
|||
import trimesh |
|||
import numpy as np |
|||
from brep2sdf.data.sampler import sample_zero_surface_points_and_normals |
|||
from brep2sdf.networks.network import gradient |
|||
import torch |
|||
import logging |
|||
|
|||
# 配置日志记录 |
|||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') |
|||
|
|||
def position_loss(pred_sdfs: torch.Tensor) -> torch.Tensor: |
|||
"""位置损失函数""" |
|||
# 保持梯度流 |
|||
squared_diff = torch.pow(pred_sdfs, 2) |
|||
return torch.mean(squared_diff) |
|||
import argparse |
|||
from skimage import measure |
|||
import time |
|||
import trimesh |
|||
import pickle |
|||
from trimesh.proximity import ProximityQuery, signed_distance, closest_point |
|||
from brep2sdf.utils.logger import logger |
|||
|
|||
def average_normal_error(normals1: torch.Tensor, normals2: torch.Tensor) -> torch.Tensor: |
|||
def sample_from_obj(obj_file, num_samples): |
|||
""" |
|||
计算平均法向量误差 (NAE) |
|||
:param normals1: 形状为 (B, 3) 的法向量张量 |
|||
:param normals2: 形状为 (B, 3) 的法向量张量 |
|||
:return: NAE 值 |
|||
从OBJ文件中采样点 |
|||
:param obj_file: OBJ文件路径 |
|||
:param num_samples: 采样点数量 |
|||
:return: 采样点数组 (N, 3) |
|||
""" |
|||
dot_products = torch.sum(normals1 * normals2, dim=-1) |
|||
absolute_dot_products = torch.abs(dot_products) |
|||
angle_errors = 1 - absolute_dot_products |
|||
return torch.mean(angle_errors) |
|||
mesh = trimesh.load(obj_file) |
|||
points, _ = mesh.sample(num_samples, return_index=True) |
|||
return points |
|||
|
|||
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: |
|||
model = torch.jit.load(model_path).to(device) |
|||
logging.info(f"成功加载模型: {model_path}") |
|||
return model |
|||
except Exception as e: |
|||
logging.error(f"加载模型 {model_path} 时出错: {e}") |
|||
return None |
|||
# Step 1: Compute signed distance |
|||
sdf = signed_distance(trimesh_mesh_ncs, points) |
|||
print(sdf) |
|||
|
|||
# Step 2: Compute normals via closest face |
|||
_, _, face_indices = closest_point(trimesh_mesh_ncs, points) |
|||
normals = trimesh_mesh_ncs.face_normals[face_indices] |
|||
|
|||
#========== |
|||
def nh(model_path, points): |
|||
model = load_model(model_path) |
|||
if model is None: |
|||
return None |
|||
try: |
|||
return model(points) |
|||
except Exception as e: |
|||
logging.error(f"调用 NH 模型时出错: {e}") |
|||
return None |
|||
# Step 3: Concatenate into (N, 7) |
|||
data = np.hstack([points, normals, sdf[:, None]]) |
|||
return data |
|||
|
|||
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: |
|||
logging.error(f"调用 mine 模型时出错: {e}") |
|||
print(f"Error computing SDF or normals: {e}") |
|||
return None |
|||
|
|||
def main(): |
|||
# 替换为实际的 obj 文件路径 |
|||
obj_file_path = "/home/wch/brep2sdf/data/gt_mesh/00000031.obj" |
|||
model_path = "/home/wch/brep2sdf/data/output_data/00000031.pt" |
|||
nh_model = "/home/wch/NH-Rep/data/output_data/00000031_0_50k_model_h.pt" |
|||
def test2(obj_file): |
|||
mesh = trimesh.load(obj_file) |
|||
surfs = sample_grid(mesh) |
|||
points = surfs[:, :3] # 提取点坐标 |
|||
normals = surfs[:, 3:6] # 提取法向量 |
|||
sdf_values = surfs[:, 6] # 提取SDF值 |
|||
|
|||
try: |
|||
# 读取 obj 文件 |
|||
mesh = trimesh.load_mesh(obj_file_path) |
|||
logging.info(f"成功读取 OBJ 文件: {obj_file_path}") |
|||
except Exception as e: |
|||
logging.error(f"读取 OBJ 文件 {obj_file_path} 时出错: {e}") |
|||
return |
|||
# 将点坐标和SDF值转换为网格格式 |
|||
grid_size = int(np.cbrt(len(points))) # 假设采样点是立方体网格 |
|||
sdf_grid = sdf_values.reshape((grid_size, grid_size, grid_size)) |
|||
|
|||
try: |
|||
# 调用采样函数 |
|||
result1 = sample_zero_surface_points_and_normals(mesh, num_samples=4096) |
|||
if result1 is None: |
|||
logging.error("采样失败，返回 None") |
|||
return |
|||
# 提取前 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") |
|||
# 使用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)}") |
|||
|
|||
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__": |
|||
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