import os import argparse import io import pickle import tarfile import numpy as np import torch import torch.nn as nn import trimesh import zstandard from scipy.spatial import KDTree from tqdm import tqdm import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from models.auto import StackedTrajectorySpatialTextRefiner 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, min_frames=16): 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 not member.isfile(): continue f = tar.extractfile(member) if f is None: continue data = pickle.load(io.BytesIO(f.read())) if "vertices" in data and "faces" in data and data["vertices"].shape[0] >= min_frames: data_list.append(data) return data_list def save_mesh_as_png(mesh, save_path): fig = plt.figure(figsize=(6, 6)) ax = fig.add_subplot(111, projection="3d") vertices = np.asarray(mesh.vertices) faces = np.asarray(mesh.faces) ax.plot_trisurf( vertices[:, 0], vertices[:, 1], vertices[:, 2], triangles=faces, cmap="viridis", edgecolor="none", alpha=0.85, ) 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 center = vertices.mean(axis=0) ax.set_xlim(center[0] - max_range, center[0] + max_range) ax.set_ylim(center[1] - max_range, center[1] + max_range) ax.set_zlim(center[2] - max_range, center[2] + max_range) ax.view_init(elev=20, azim=45) plt.axis("off") plt.savefig(save_path, bbox_inches="tight", pad_inches=0, dpi=150) plt.close(fig) def to_tensor(x, device, dtype=None): if isinstance(x, torch.Tensor): x = x.detach().clone() return x.to(device=device, dtype=dtype) if dtype is not None else x.to(device=device) return torch.as_tensor(x, device=device, dtype=dtype) @torch.no_grad() def encode_sequence_latents(vae, vertices_seq, faces, device, num_frames, num_samples): latents = [] f_t = to_tensor(faces, device=device, dtype=torch.float32) for frame_idx in range(num_frames): v_t = to_tensor(vertices_seq[frame_idx], device=device, dtype=torch.float32) latent = vae.encode_mesh(v_t, f_t, num_samples=num_samples) if isinstance(latent, np.ndarray): latent = torch.from_numpy(latent).to(device) latents.append(latent.squeeze(0)) return torch.stack(latents, dim=0) @torch.no_grad() def autoregressive_predict(auto_model, initial_latent, text_embed, num_future_frames, num_tokens, device): current_context = initial_latent.unsqueeze(0).unsqueeze(0).to(device) shortest_path_matrix = torch.zeros((1, num_tokens, num_tokens), device=device, dtype=torch.long) pred_latents = [current_context[:, 0:1]] for _ in range(num_future_frames): pred_seq = auto_model( current_context, text_embed=text_embed, shortest_path_matrix=shortest_path_matrix, ) next_latent = pred_seq[:, -1:, :, :] pred_latents.append(next_latent) current_context = torch.cat([current_context, next_latent], dim=1) return pred_latents def compute_mesh_metrics(mesh_recon, mesh_gt, num_points=2048, tau=0.05): pts_gt, _ = trimesh.sample.sample_surface(mesh_gt, num_points) pts_recon, _ = trimesh.sample.sample_surface(mesh_recon, num_points) 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 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 return d_r2r, d_raw2r, cd, f_score def build_model(args, device): return StackedTrajectorySpatialTextRefiner( num_groups=args.num_refine_groups, temporal_kwargs={ "num_tokens": args.num_tokens, "token_dim": args.token_dim, "max_seq_len": args.max_seq_len, "num_blocks": args.temporal_blocks, "compress_layers": args.compress_layers, }, spatial_kwargs={ "num_tokens": args.num_tokens, "token_dim": args.token_dim, "num_blocks": args.spatial_blocks, }, text_kwargs={ "token_dim": args.token_dim, "text_dim": args.token_dim, }, ).to(device) def load_checkpoint(auto_model, text_proj, checkpoint_path): ckpt = torch.load(checkpoint_path, map_location="cpu") if "auto" not in ckpt or "proj" not in ckpt: raise KeyError("Checkpoint must contain keys 'auto' and 'proj'.") auto_model.load_state_dict(ckpt["auto"]) text_proj.load_state_dict(ckpt["proj"]) return ckpt def evaluate(args): model_device = torch.device(args.model_device) vae_device = torch.device(args.vae_device) clip_device = torch.device(args.clip_device) os.makedirs(args.output_dir, exist_ok=True) vae = TripoSGVaeWrapper(device=str(vae_device)) vae.eval() text_encoder = FrozenCLIPTextEncoder(args.clip_model_path, device=str(clip_device)) text_proj = TextProjectionLayer(in_dim=512, out_dim=args.token_dim).to(clip_device) auto_model = build_model(args, model_device) load_checkpoint(auto_model, text_proj, args.checkpoint_path) auto_model.eval() text_proj.eval() dataset = get_full_trajectories(args.archive_path, limit=args.limit, min_frames=args.num_frames) metrics = {"mae_r2r": [], "mae_raw2r": [], "cd": [], "fscore": []} print(f"Loaded {len(dataset)} trajectories.") print(f"Model device: {model_device}, VAE device: {vae_device}, CLIP device: {clip_device}") for traj_idx, traj_data in enumerate(tqdm(dataset)): if args.visualize_every > 0 and traj_idx % args.visualize_every != 0: continue traj_vis_dir = os.path.join(args.output_dir, f"{traj_idx:04d}") os.makedirs(traj_vis_dir, exist_ok=True) vertices_seq = traj_data["vertices"][: args.num_frames] faces = traj_data["faces"] caption = traj_data.get("caption", "") gt_latents = encode_sequence_latents( vae=vae, vertices_seq=vertices_seq, faces=faces, device=vae_device, num_frames=args.num_frames, num_samples=args.num_surface_samples, ) text_tokens = text_encoder([caption]) text_embed = text_proj(text_tokens.to(clip_device)).to(model_device) pred_latents = autoregressive_predict( auto_model=auto_model, initial_latent=gt_latents[0].to(model_device), text_embed=text_embed, num_future_frames=args.num_frames - 1, num_tokens=args.num_tokens, device=model_device, ) 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, "000_gt.png")) for frame_idx in range(1, args.num_frames): latent_t = pred_latents[frame_idx].to(vae_device).squeeze(1) geometric_func = lambda x, latent=latent_t: vae.decode_latent(latent, sampled_points=x) output_geometry = hierarchical_extract_geometry( geometric_func, device=str(vae_device), 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_gt = trimesh.Trimesh(vertices=vertices_seq[frame_idx], faces=faces, process=False) save_mesh_as_png(mesh_recon, os.path.join(traj_vis_dir, f"{frame_idx:03d}_pred.png")) if args.save_gt_frames: save_mesh_as_png(mesh_gt, os.path.join(traj_vis_dir, f"{frame_idx:03d}_gt.png")) d_r2r, d_raw2r, cd, f_score = compute_mesh_metrics( mesh_recon, mesh_gt, num_points=args.metric_points, tau=args.fscore_tau, ) metrics["mae_r2r"].append(d_r2r) metrics["mae_raw2r"].append(d_raw2r) metrics["cd"].append(cd) metrics["fscore"].append(f_score) running = {k: float(np.mean(v)) for k, v in metrics.items() if len(v) > 0} print(f"Trajectory {traj_idx:04d} | {running}") print("\n--- Evaluation Results ---") for key, values in metrics.items(): print(f"{key}: {np.mean(values):.6f}" if values else f"{key}: no samples") def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--checkpoint_path", type=str, required=True) parser.add_argument("--archive_path", type=str, required=True) parser.add_argument("--clip_model_path", type=str, default="/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/pretrain/clip-vit-base-patch32") parser.add_argument("--output_dir", type=str, default="/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/mesh_vis_results") parser.add_argument("--num_tokens", type=int, default=512) parser.add_argument("--token_dim", type=int, default=64) parser.add_argument("--num_refine_groups", type=int, default=3) parser.add_argument("--temporal_blocks", type=int, default=6) parser.add_argument("--compress_layers", type=int, default=3) parser.add_argument("--spatial_blocks", type=int, default=3) parser.add_argument("--max_seq_len", type=int, default=16) parser.add_argument("--num_frames", type=int, default=16) parser.add_argument("--limit", type=int, default=400) parser.add_argument("--visualize_every", type=int, default=74) parser.add_argument("--num_surface_samples", type=int, default=51200) parser.add_argument("--metric_points", type=int, default=2048) parser.add_argument("--fscore_tau", type=float, default=0.05) parser.add_argument("--save_gt_frames", action="store_true") parser.add_argument("--model_device", type=str, default="cuda:0") parser.add_argument("--vae_device", type=str, default="cuda:0") parser.add_argument("--clip_device", type=str, default="cuda:0") return parser.parse_args() if __name__ == "__main__": evaluate(parse_args())