test_v2.py

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()