File size: 3,066 Bytes
c25da31 | 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 | import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import argparse
import yaml
import torch
from data.smal_quadrupeds import SMALQuadrupedGenerator, FAMILIES
from models.flow_matching import OptimalTransportFlowMatching
from pipeline.metrics import compute_generation_metrics, surface_distance
from scripts.evaluate_families import load_model, add_sampling_args, sampling_kwargs
def main():
"""
Standard generation metrics (MMD, COV, 1-NNA with Chamfer Distance): creatures generated from unseen
skeletons against as many real creatures drawn independently from the same families.
"""
parser = argparse.ArgumentParser(description="Evaluate generation quality with MMD, COV and 1-NNA.")
parser.add_argument("--config", type=str, default="configs/default_config.yaml")
parser.add_argument("--checkpoint", type=str, default=None)
parser.add_argument("--samples-per-family", type=int, default=40)
parser.add_argument("--batch-size", type=int, default=20, help="skeletons sampled together")
add_sampling_args(parser)
args = parser.parse_args()
with open(args.config, "r", encoding="utf-8") as f:
cfg = yaml.safe_load(f)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
ckpt_path = args.checkpoint or os.path.join(cfg["training"]["checkpoint_dir"], f"model_epoch_{cfg['training']['num_epochs']:03d}.pt")
fm = OptimalTransportFlowMatching(load_model(cfg, ckpt_path, device, args.raw_weights))
sample_kwargs = sampling_kwargs(args, cfg, device)
# Unseen seed, disjoint from the training set (seed 42); references are separate draws of the same families
torch.manual_seed(0)
generator = SMALQuadrupedGenerator(seed=123)
to_tensor = lambda a: torch.from_numpy(a).to(device)
gen, ref, gaps = [], [], []
for family in FAMILIES:
skeletons = []
for _ in range(args.samples_per_family):
skeletons.append(to_tensor(generator.generate_instance(family)["joints"]))
ref.append(to_tensor(generator.generate_instance(family)["points"]))
radii = generator.bone_radii(family)
for batch in torch.stack(skeletons).split(args.batch_size):
points = fm.sample(batch, **sample_kwargs)[1]
gen.extend(points)
gaps.extend(surface_distance(p, j, radii) for p, j in zip(points, batch))
metrics = compute_generation_metrics(torch.stack(gen), torch.stack(ref))
print(f"MMD-CD: {metrics['mmd']:.5f} | COV-CD: {metrics['cov']:.0%} | 1-NNA-CD: {metrics['1nna']:.1%} (50% is ideal; "
f"generated {metrics['1nna_gen']:.0%}, real {metrics['1nna_ref']:.0%})")
# Distance to the procedural surface the skeleton defines (real shapes: median about 0.0004)
gap = torch.cat(gaps)
print(f"Surface distance: median {gap.abs().median():.4f} | p90 {gap.abs().quantile(0.9):.4f} | "
f"mean signed {gap.mean():+.4f} (positive = outside) | {len(gen)} vs {len(ref)} shapes")
if __name__ == "__main__":
main()
|