import io import os import tarfile import pickle import zstandard import matplotlib.pyplot as plt import numpy as np import torch import trimesh # 引入刚才写好的 VAE 包装类 from models.vae import TripoSGVaeWrapper # 引入 TripoSG 官方提供的层次化几何提取工具 from triposg.inference_utils import hierarchical_extract_geometry # 配置路径 archive_path = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/DyMesh_50000v_16f_0001_part_00" triposg_weights = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/TripoSG/pretrained_weights/TripoSG" output_image_path = "vae_reconstruction_test.png" def get_first_mesh_frame(archive_file): """从数据集中提取第一个有效物体的第0帧""" with open(archive_file, 'rb') as fh: dctx = zstandard.ZstdDecompressor() with dctx.stream_reader(fh) as reader: with tarfile.open(fileobj=reader, mode='r|') as tar: for member in tar: if member.isfile(): f = tar.extractfile(member) if f is not None: data = pickle.load(io.BytesIO(f.read())) if isinstance(data, dict) and 'vertices' in data and 'faces' in data: if data['vertices'].shape[0] == 16: # 提取第0帧 (形状为 [V, 3]) 和 共享的面片 v_f0 = data['vertices'][0] faces = data['faces'] short_name = os.path.basename(member.name).split('.')[0] return v_f0, faces, short_name raise RuntimeError("未在压缩包中找到合法的 16帧 动态 Mesh 数据。") if __name__ == "__main__": device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.float16 # 1. 读取原数据并抽取静态帧 print("📦 正在从数据包中读取测试 Mesh...") raw_vertices, raw_faces, obj_name = get_first_mesh_frame(archive_path) print(f"成功提取物体 [{obj_name}] 的第一帧. 顶点数: {raw_vertices.shape[0]}, 面数: {raw_faces.shape[0]}") # 2. 对原网格在内存中做归一化，以便后续做对齐可视化 mesh_raw = trimesh.Trimesh(vertices=raw_vertices, faces=raw_faces, process=False) max_edge = np.max(mesh_raw.bounding_box.extents) mesh_raw.vertices = (mesh_raw.vertices - mesh_raw.center_mass) * (2.0 / max_edge) # 3. 初始化 VAE 包装模型 print("🔮 正在初始化冻结的 TripoSG VAE 模型...") vae_wrapper = TripoSGVaeWrapper(weights_dir=triposg_weights, device=device, dtype=dtype) # 4. 编码 (Mesh -> Latent Code) # 修改 3. 编码部分的调用逻辑： print("🚀 正在将 Mesh 表面密集点云编码至隐空间...") # 将 numpy 数组显式转换为 torch.Tensor v_tensor = torch.from_numpy(raw_vertices) f_tensor = torch.from_numpy(raw_faces) latent_code = vae_wrapper.encode_mesh(v_tensor, f_tensor, num_samples=204800) print(f"编码成功！隐特征张量形状: {latent_code.shape}") # 5. 解码重建 (Latent Code -> SDF -> Marching Cubes Mesh) print("📐 正在通过层次化八叉树提取器（Hierarchical Octree）解码并重建 Mesh...") geometric_func = lambda x: vae_wrapper.decode_latent(latent_code, sampled_points=x) output_geometry = hierarchical_extract_geometry( geometric_func, device=device, bounds=(-1.005, -1.005, -1.005, 1.005, 1.005, 1.005), dense_octree_depth=8, # 密度深度 hierarchical_octree_depth=9, # 恢复细节的八叉树总深度 ) # 将提取出的解算数据转换为 trimesh 对象 recon_v, recon_f = output_geometry[0] mesh_recon = trimesh.Trimesh(vertices=recon_v.astype(np.float32), faces=recon_f, process=False) print(f"重建成功！重建后 Mesh 顶点数: {mesh_recon.vertices.shape[0]}, 面数: {mesh_recon.faces.shape[0]}") # 🌟 5.5 进行几何指标定量评估 (Evaluation Metric) - 使用 Scipy KDTree 修复版本 print("📊 正在执行 KDTree 算法，评测重建精度指标...") from scipy.spatial import KDTree num_eval_samples = 50000 # 在两个几何表面分别均匀采样 5万 个测试点 pts_raw, _ = trimesh.sample.sample_surface(mesh_raw, num_eval_samples) pts_recon, _ = trimesh.sample.sample_surface(mesh_recon, num_eval_samples) # 建立空间平衡二叉树，用于秒级检索最近邻点 tree_raw = KDTree(pts_raw) tree_recon = KDTree(pts_recon) # 1. 计算重建点到真值点云的最近距离 (Recon -> Raw) -> 代表重建偏离误差 dists_recon_to_raw, _ = tree_raw.query(pts_recon, k=1) # 2. 计算真值点到重建点云的最近距离 (Raw -> Recon) -> 代表细节遗漏误差 dists_raw_to_recon, _ = tree_recon.query(pts_raw, k=1) # 计算均值 mae_recon_to_raw = np.mean(dists_recon_to_raw) mae_raw_to_recon = np.mean(dists_raw_to_recon) # 模拟双向倒角距离 (Chamfer Distance) chamfer_dist_approx = mae_recon_to_raw + mae_raw_to_recon # 计算在 1% 空间尺度阈值下的 F-Score obj_scale = np.linalg.norm(mesh_raw.bounding_box.extents) tau = 0.01 * obj_scale # 阈值设置为物体对角线长度的 1% precision = np.mean(dists_recon_to_raw < tau) recall = np.mean(dists_raw_to_recon < tau) f_score = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0 # 打印评估报告 print("\n" + "="*50) print("📈 VAE Mesh 自编码重建评测报告 (KDTree 加速版):") print("="*50) print(f"1. 重建偏离误差 (Recon -> Raw MAE): {mae_recon_to_raw:.6f} (越小越好)") print(f"2. 细节遗漏误差 (Raw -> Recon MAE): {mae_raw_to_recon:.6f} (越小越好)") print(f"3. 倒角距离指标 (Chamfer Distance): {chamfer_dist_approx:.6f} (越小越好)") print(f"4. 表面对齐分数 (F-Score @ 1%): {f_score * 100:.2f}% (越大越好)") print("="*50 + "\n") # 6. 使用 Matplotlib 绘制同视角对比图 print("🎨 正在绘制原网格与 VAE 重建网格的对比效果图...") fig = plt.figure(figsize=(14, 6)) # 共同视角参数控制，防止发生视差偏离 all_points = np.concatenate([mesh_raw.vertices, mesh_recon.vertices], axis=0) mid_x, mid_y, mid_z = (all_points.max(axis=0) + all_points.min(axis=0)) / 2.0 max_range = (all_points.max(axis=0) - all_points.min(axis=0)).max() / 2.0 # 子图 1：原始 Mesh（归一化后） ax1 = fig.add_subplot(1, 2, 1, projection='3d') ax1.plot_trisurf(mesh_raw.vertices[:, 0], mesh_raw.vertices[:, 1], mesh_raw.vertices[:, 2], triangles=mesh_raw.faces, cmap='coolwarm', edgecolor='none', alpha=0.7) ax1.set_title(f"Original Mesh (Normalized)\n({obj_name})", fontsize=12, weight='bold') ax1.set_xlim(mid_x - max_range, mid_x + max_range) ax1.set_ylim(mid_y - max_range, mid_y + max_range) ax1.set_zlim(mid_z - max_range, mid_z + max_range) ax1.axis('off') ax1.view_init(elev=20, azim=45) # 设定固定视角 # 子图 2：VAE 解码重建的 Mesh ax2 = fig.add_subplot(1, 2, 2, projection='3d') ax2.plot_trisurf(mesh_recon.vertices[:, 0], mesh_recon.vertices[:, 1], mesh_recon.vertices[:, 2], triangles=mesh_recon.faces, cmap='viridis', edgecolor='none', alpha=0.7) # 在标题上动态附带上 F-Score 信息，方便肉眼和指标联动对照 ax2.set_title(f"Reconstructed Mesh\n(F-Score: {f_score * 100:.2f}%)", fontsize=12, weight='bold') ax2.set_xlim(mid_x - max_range, mid_x + max_range) ax2.set_ylim(mid_y - max_range, mid_y + max_range) ax2.set_zlim(mid_z - max_range, mid_z + max_range) ax2.axis('off') ax2.view_init(elev=20, azim=45) # 设定相同视角 plt.tight_layout() plt.savefig(output_image_path, dpi=150, bbox_inches='tight') plt.close() print(f"💾 对比可视化结果已成功保存至: {output_image_path}")