import os import torch import torch.nn as nn import trimesh import numpy as np import io import tarfile import pickle import zstandard from scipy.spatial import KDTree from tqdm import tqdm import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D # 引入你的模型组件 from models.auto import TemporalTrajectoryPredictor, SpatialGraphRefiner from models.vae import TripoSGVaeWrapper from models.clip import FrozenCLIPTextEncoder from triposg.inference_utils import hierarchical_extract_geometry # 辅助类 class TextProjectionLayer(nn.Module): def __init__(self, in_dim=512, out_dim=64): super().__init__() self.net = nn.Sequential(nn.Linear(in_dim, in_dim), nn.ReLU(), nn.Linear(in_dim, out_dim)) self.residual = nn.Linear(in_dim, out_dim) if in_dim != out_dim else nn.Identity() def forward(self, x): return self.net(x) + self.residual(x) def get_full_trajectories(archive_file, limit=20): """提取完整的轨迹数据及caption""" data_list = [] 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 len(data_list) >= limit: break if member.isfile(): f = tar.extractfile(member) if f: data = pickle.load(io.BytesIO(f.read())) # 确保数据中包含 caption if 'vertices' in data and data['vertices'].shape[0] >= 16: data_list.append(data) return data_list def save_mesh_as_png(mesh, save_path): """使用matplotlib将mesh渲染并保存为png图片，兼容无显示器的服务器环境""" fig = plt.figure(figsize=(6, 6)) ax = fig.add_subplot(111, projection='3d') # 提取顶点和面数据 vertices = mesh.vertices faces = mesh.faces # 绘制3D三角网格 ax.plot_trisurf(vertices[:, 0], vertices[:, 1], vertices[:, 2], triangles=faces, cmap='viridis', edgecolor='none', alpha=0.8) # 保持各轴比例统一并隐藏坐标轴以使视觉效果更干净 ax.set_xlabel('X') ax.set_ylabel('Y') ax.set_zlabel('Z') # 设置一个较好的初始观察视角 ax.view_init(elev=20, azim=45) # 强制使轴范围一致 max_range = np.array([vertices[:, 0].max()-vertices[:, 0].min(), vertices[:, 1].max()-vertices[:, 1].min(), vertices[:, 2].max()-vertices[:, 2].min()]).max() / 2.0 mid_x = (vertices[:, 0].max()+vertices[:, 0].min()) * 0.5 mid_y = (vertices[:, 1].max()+vertices[:, 1].min()) * 0.5 mid_z = (vertices[:, 2].max()+vertices[:, 2].min()) * 0.5 ax.set_xlim(mid_x - max_range, mid_x + max_range) ax.set_ylim(mid_y - max_range, mid_y + max_range) ax.set_zlim(mid_z - max_range, mid_z + max_range) plt.axis('off') plt.savefig(save_path, bbox_inches='tight', pad_inches=0, dpi=150) plt.close(fig) def evaluate(): # 检测可用GPU数量并进行针对性非对称分配 num_gpus = torch.cuda.device_count() # 主卡分配给显存最多的 cuda:0 dev_main = "cuda:0" if num_gpus >= 3: # 显存大户（VAE和自回归Predictor）全部塞给大显存的 cuda:0 dev_vae = dev_main dev_temporal = dev_main # 轻量级的文本端和空间端移出主卡，分流到剩余卡上 dev_clip = "cuda:1" dev_spatial = "cuda:2" print(f"检测到 {num_gpus} 张卡。策略：VAE与Temporal驻留大显存卡({dev_main})，CLIP分流至(cuda:1)，Spatial分流至(cuda:2)。") elif num_gpus == 2: dev_vae = dev_main dev_temporal = dev_main dev_clip = "cuda:1" dev_spatial = "cuda:1" print(f"检测到 2 张卡。策略：VAE与Temporal驻留大显存卡({dev_main})，CLIP与Spatial分流至(cuda:1)。") else: dev_vae = dev_clip = dev_temporal = dev_spatial = dev_main print(f"仅有单卡或策略回退，所有模型将挤在 {dev_main} 上运行。") checkpoint_path = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/checkpoints/20260608_011428/model_step_1500.pth" archive_path = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/shards/shard_0000.tar.zst" clip_path = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/pretrain/clip-vit-base-patch32" vis_base_dir = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/mesh_vis_results" # 1. 初始化模型并发送到指定卡上 vae = TripoSGVaeWrapper(device=dev_vae) vae.eval() text_encoder = FrozenCLIPTextEncoder(clip_path, device=dev_clip) temporal_model = TemporalTrajectoryPredictor(num_tokens=512, token_dim=64, max_seq_len=16).to(dev_temporal) spatial_model = SpatialGraphRefiner(num_tokens=512, token_dim=64, text_dim=64).to(dev_spatial) text_proj = TextProjectionLayer(in_dim=512, out_dim=64).to(dev_clip) # 加载权重 (先加载至CPU，避免加载时在主卡产生双倍显存突刺) ckpt = torch.load(checkpoint_path, map_location="cpu") temporal_model.load_state_dict(ckpt['temporal']) spatial_model.load_state_dict(ckpt['spatial']) text_proj.load_state_dict(ckpt['proj']) temporal_model.eval() spatial_model.eval() text_proj.eval() # 准备数据 dataset = get_full_trajectories(archive_path, limit=400) metrics = {"mae_r2r": [], "mae_raw2r": [], "cd": [], "fscore": []} print(f"开始测试，共 {len(dataset)} 条轨迹...") for traj_idx, traj_data in enumerate(tqdm(dataset)): if traj_idx % 25 == 0 and traj_idx >= 1: # 创建输出文件夹 traj_vis_dir = os.path.join(vis_base_dir, f"{traj_idx:04d}") os.makedirs(traj_vis_dir, exist_ok=True) vertices_seq = traj_data['vertices'][:16] faces = traj_data['faces'] caption = traj_data.get('caption') # 1. 编码所有GT帧到 Latent 以及处理文本特征 with torch.no_grad(): gt_latents = [] for t in range(16): v_t = torch.from_numpy(vertices_seq[t]).to(device=dev_vae, dtype=torch.float32) f_t = torch.from_numpy(faces).to(device=dev_vae, dtype=torch.float32) l = vae.encode_mesh(v_t, f_t) gt_latents.append(l.squeeze(0)) gt_latents = torch.stack(gt_latents, dim=0) # [16, 512, 64] (当前在 dev_vae 即 cuda:0) # 映射文本特征 (在指定的 dev_clip 卡上完成) text_embed = text_proj(text_encoder([caption])).unsqueeze(1) # [1, 1, 64] # 2. 自回归预测 (进行精细的跨卡张量传递) current_context = gt_latents[0:1].unsqueeze(0).to(dev_temporal) # [1, 1, 512, 64] # 保存用作后面解码的预测隐变量列表 (统一放在负责解码的 dev_vae 也就是 cuda:0 上) pred_latents = [gt_latents[0:1].unsqueeze(0).to(dev_vae)] for t in range(15): # 步骤 A: 时间序列模型自回归预测 (在大显存卡 dev_temporal 执行) text_embed_to_spatial = text_embed.to(dev_temporal) next_t = temporal_model(current_context, text_embed_to_spatial) # [1, 1, 512, 64] next_t = next_t[:, -1:, :, :] # 步骤 B: 传递数据到分流卡进行空间细化 (转移至 dev_spatial) text_embed_to_spatial = text_embed.to(dev_spatial) next_t_to_spatial = next_t.to(dev_spatial) spatial_mask = torch.zeros((1, 512, 512), device=dev_spatial) next_t_refined = spatial_model( next_t_to_spatial, text_embed_to_spatial, spatial_mask ) next_t_refined = next_t_refined[:, -1:, :, :] # 步骤 C: 将精细化后的隐变量送回负责解码的卡上备份 pred_latents.append(next_t_refined.to(dev_vae)) # 同时更新自回归上下文 (送回 dev_temporal 卡) next_t_to_temporal = next_t_refined.to(dev_temporal) current_context = torch.cat([current_context, next_t_to_temporal], dim=1) # 3. 处理第 0 帧 (GT 帧) 的可视化 with torch.no_grad(): mesh_gt_0 = trimesh.Trimesh(vertices=vertices_seq[0], faces=faces, process=False) save_mesh_as_png(mesh_gt_0, os.path.join(traj_vis_dir, f"{0:03d}.png")) # 4. 解码并计算指标 (在 dev_vae 即大显存卡 cuda:0 上运行) for t in range(1, 16): latent_t = pred_latents[t].squeeze(1).to(dev_vae) geometric_func = lambda x: vae.decode_latent(latent_t, sampled_points=x) output_geometry = hierarchical_extract_geometry( geometric_func, device=dev_vae, bounds=(-1.005, -1.005, -1.005, 1.005, 1.005, 1.005) ) recon_v, recon_f = output_geometry[0] mesh_recon = trimesh.Trimesh(vertices=recon_v, faces=recon_f, process=False) # 可视化并保存当前生成的 mesh 帧图片 save_mesh_as_png(mesh_recon, os.path.join(traj_vis_dir, f"{t:03d}.png")) mesh_gt = trimesh.Trimesh(vertices=vertices_seq[t], faces=faces, process=False) pts_gt, _ = trimesh.sample.sample_surface(mesh_gt, 2048) pts_recon, _ = trimesh.sample.sample_surface(mesh_recon, 2048) tree_gt = KDTree(pts_gt) tree_recon = KDTree(pts_recon) d_r2r = np.mean(tree_gt.query(pts_recon, k=1)[0]) d_raw2r = np.mean(tree_recon.query(pts_gt, k=1)[0]) cd = d_r2r + d_raw2r tau = 0.05 precision = np.mean(tree_gt.query(pts_recon, k=1)[0] < tau) recall = np.mean(tree_recon.query(pts_gt, k=1)[0] < tau) f_score = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0 metrics["mae_r2r"].append(d_r2r) metrics["mae_raw2r"].append(d_raw2r) metrics["cd"].append(cd) metrics["fscore"].append(f_score) for k, v in metrics.items(): print(f"{k}: {np.mean(v):.6f}") print("\n--- 测试结果 ---") for k, v in metrics.items(): print(f"{k}: {np.mean(v):.6f}") if __name__ == "__main__": evaluate()