File size: 4,694 Bytes
f0fc238 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | import io
import os
import tarfile
import pickle
import zstandard
import numpy as np
import torch
import trimesh
from scipy.spatial import KDTree
# 引入 VAE 包装类和几何提取工具
from models.vae import TripoSGVaeWrapper
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"
def get_mesh_frames(archive_file, limit=100):
"""从数据集中提取前 N 个有效物体的第0帧"""
extracted_data = []
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(extracted_data) >= limit:
break
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] >= 1:
v_f0 = data['vertices'][0]
faces = data['faces']
short_name = os.path.basename(member.name).split('.')[0]
extracted_data.append((v_f0, faces, short_name))
return extracted_data
if __name__ == "__main__":
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16
print(f"📦 正在从数据包中读取前 100 个测试 Mesh...")
mesh_list = get_mesh_frames(archive_path, limit=100)
# 初始化 VAE
print("🔮 正在初始化冻结的 TripoSG VAE 模型...")
vae_wrapper = TripoSGVaeWrapper(weights_dir=triposg_weights, device=device, dtype=dtype)
metrics_sum = {"mae_r2r": 0.0, "mae_raw2r": 0.0, "cd": 0.0, "fscore": 0.0}
count = 0
print(f"\n{'ID':<5} | {'Recon->Raw':<12} | {'Raw->Recon':<12} | {'CD':<10} | {'F-Score':<8}")
print("-" * 65)
for i, (raw_vertices, raw_faces, obj_name) in enumerate(mesh_list):
# 预处理
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)
# 编码与解码
v_tensor = torch.from_numpy(mesh_raw.vertices).to(device=device, dtype=dtype)
f_tensor = torch.from_numpy(raw_faces).to(device=device, dtype=torch.int32)
latent_code = vae_wrapper.encode_mesh(v_tensor, f_tensor, num_samples=204800)
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
)
recon_v, recon_f = output_geometry[0]
mesh_recon = trimesh.Trimesh(vertices=recon_v.astype(np.float32), faces=recon_f, process=False)
# 评估
pts_raw, _ = trimesh.sample.sample_surface(mesh_raw, 204800)
pts_recon, _ = trimesh.sample.sample_surface(mesh_recon, 204800)
tree_raw = KDTree(pts_raw)
tree_recon = KDTree(pts_recon)
d_r2r = np.mean(tree_raw.query(pts_recon, k=1)[0])
d_raw2r = np.mean(tree_recon.query(pts_raw, k=1)[0])
cd = d_r2r + d_raw2r
tau = 0.01 * np.linalg.norm(mesh_raw.bounding_box.extents)
precision = np.mean(tree_raw.query(pts_recon, k=1)[0] < tau)
recall = np.mean(tree_recon.query(pts_raw, k=1)[0] < tau)
f_score = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0
# 记录
metrics_sum["mae_r2r"] += d_r2r
metrics_sum["mae_raw2r"] += d_raw2r
metrics_sum["cd"] += cd
metrics_sum["fscore"] += f_score
count += 1
print(f"{i+1:<5} | {d_r2r:.4f} | {d_raw2r:.4f} | {cd:.4f} | {f_score*100:.2f}%")
# 输出平均值
if count > 0:
print("-" * 65)
print(f"平均值: | {metrics_sum['mae_r2r']/count:.4f} | {metrics_sum['mae_raw2r']/count:.4f} | {metrics_sum['cd']/count:.4f} | {(metrics_sum['fscore']/count)*100:.2f}%") |