File size: 2,416 Bytes
7ca3cc4 38b41a7 7ca3cc4 | 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 | import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import yaml
import torch
from data.dataset import CreatureDataset
from models.point_dit import PointDiT
from models.flow_matching import OptimalTransportFlowMatching
from pipeline.metrics import compute_chamfer_distance, compute_skeleton_coverage
def main():
with open("configs/default_config.yaml", "r", encoding="utf-8") as f:
cfg = yaml.safe_load(f)
dataset = CreatureDataset(num_samples_per_family=5, num_points=1024, augment=False, cache_dir=None)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = PointDiT(
point_dim=cfg["model"]["point_dim"],
hidden_dim=cfg["model"]["hidden_dim"],
num_layers=cfg["model"]["num_layers"],
num_heads=cfg["model"]["num_heads"],
num_joints=cfg["model"]["num_joints"],
spatial_distance_scale=cfg["model"]["spatial_distance_scale"],
fourier_octaves=cfg["model"].get("fourier_octaves", 0),
local_k=cfg["model"].get("local_k", 0)
)
ckpt_path = os.path.join(cfg["training"]["checkpoint_dir"], f"model_epoch_{cfg['training']['num_epochs']:03d}.pt")
if os.path.exists(ckpt_path):
ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)
model.load_state_dict(ckpt.get("ema_state", ckpt["model_state"]))
print(f"Loaded checkpoint from {ckpt_path}")
else:
print(f"Warning: no checkpoint at {ckpt_path}, evaluating an untrained model.")
model.to(device)
fm = OptimalTransportFlowMatching(model)
print("Evaluating test metrics on generated creatures...")
all_cd, all_cov = [], []
for i in range(min(5, len(dataset))):
sample = dataset[i]
joints = sample["joints"].unsqueeze(0).to(device)
target_pts = sample["points"].unsqueeze(0).to(device)
_, gen_pts = fm.sample(joints, num_points=1024, num_steps=15)
cd = compute_chamfer_distance(gen_pts, target_pts).item()
cov = compute_skeleton_coverage(joints, gen_pts).item()
all_cd.append(cd)
all_cov.append(cov)
print(f"Mean Chamfer Distance: {sum(all_cd)/len(all_cd):.5f}")
print(f"Mean Skeleton Coverage: {sum(all_cov)/len(all_cov):.5f}")
if __name__ == "__main__":
main()
|