test_vae.py

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}")