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