File size: 5,130 Bytes
4c3d957 cca6827 4c3d957 a3752bd 4c3d957 cca6827 4c3d957 cca6827 4c3d957 cca6827 4c3d957 | 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """Geometry metrics vs a GT mesh (LOCKED spec).
Both meshes are already canonical in [-0.5,0.5]^3; NO ICP / NO alignment.
100k surface points sampled per mesh (same count both sides) via
pytorch3d.ops.sample_points_from_meshes. NN distances via knn_points.
Reported:
* Chamfer L1 (headline): mean over both directions of the sqrt Euclidean
nearest-neighbour distance. Units = fraction of the unit-cube edge.
* Chamfer L2 : pytorch3d.loss.chamfer_distance default (mean squared).
* F-score@{0.01,0.02,0.05}: P = frac pred within tau of GT,
R = frac GT within tau of pred, F = 2PR/(P+R). Headline F@0.02.
* Normal Consistency: mean abs cosine of normals at mutual NN (both dirs).
* Volume IoU@128^3: solid voxel occupancy on a shared 128^3 grid, I/U.
"""
from __future__ import annotations
# FINAL: = metrics/appeval/geometry.py + extent-relative CD/F keys, optional vol IoU. Original keys unchanged.
import numpy as np
import torch
import trimesh
from pytorch3d.structures import Meshes
from pytorch3d.ops import sample_points_from_meshes, knn_points
from pytorch3d.loss import chamfer_distance
N_SAMPLES = 100_000
TAUS = (0.01, 0.02, 0.05)
VOX_RES = 128
def _to_p3d(mesh: trimesh.Trimesh, device):
v = torch.as_tensor(np.asarray(mesh.vertices), dtype=torch.float32, device=device)
f = torch.as_tensor(np.asarray(mesh.faces), dtype=torch.int64, device=device)
return Meshes(verts=[v], faces=[f])
def _sample(mesh, device, n=N_SAMPLES, seed=0):
torch.manual_seed(seed)
m = _to_p3d(mesh, device)
pts, nrm = sample_points_from_meshes(m, n, return_normals=True)
return pts, nrm # (1,n,3), (1,n,3)
def _voxelize_solid(mesh: trimesh.Trimesh, res=VOX_RES):
"""Solid occupancy on a fixed res^3 grid over [-0.5,0.5]^3 (bool array)."""
pitch = 1.0 / res
grid = np.zeros((res, res, res), dtype=bool)
try:
vox = mesh.voxelized(pitch=pitch)
try:
vox = vox.fill()
except Exception:
pass # non-watertight -> hollow occupancy, still comparable
pts = np.asarray(vox.points) # world-space centres of occupied cells
except Exception:
# last-resort: dense surface samples
pts, _ = trimesh.sample.sample_surface(mesh, 2_000_000, seed=0)
pts = np.asarray(pts)
idx = np.floor((pts + 0.5) * res).astype(np.int64)
idx = np.clip(idx, 0, res - 1)
grid[idx[:, 0], idx[:, 1], idx[:, 2]] = True
return grid
def geometry_metrics(pred_mesh: trimesh.Trimesh, gt_mesh: trimesh.Trimesh,
device="cuda", seed=0, with_vol=True):
"""Return a dict of all geometry metrics. NO alignment applied."""
pp, pn = _sample(pred_mesh, device, seed=seed)
gp, gn = _sample(gt_mesh, device, seed=seed)
# NN both directions (knn returns squared dists)
p2g = knn_points(pp, gp, K=1)
g2p = knn_points(gp, pp, K=1)
d_p2g = p2g.dists[..., 0].clamp_min(0).sqrt()[0] # pred->gt (n,)
d_g2p = g2p.dists[..., 0].clamp_min(0).sqrt()[0] # gt->pred (n,)
cd_l1 = 0.5 * (d_p2g.mean() + d_g2p.mean())
cd_l2, _ = chamfer_distance(pp, gp) # default mean squared
fscore = {}
for tau in TAUS:
precision = (d_p2g < tau).float().mean()
recall = (d_g2p < tau).float().mean()
f = (2 * precision * recall / (precision + recall + 1e-12))
fscore[tau] = {"precision": float(precision), "recall": float(recall),
"f": float(f)}
# normal consistency at mutual NN, both directions
idx_p2g = p2g.idx[..., 0][0] # for each pred pt, nearest gt idx
idx_g2p = g2p.idx[..., 0][0]
gn0, pn0 = gn[0], pn[0]
cos_p = (pn0 * gn0[idx_p2g]).sum(-1).abs()
cos_g = (gn0 * pn0[idx_g2p]).sum(-1).abs()
nc = 0.5 * (cos_p.mean() + cos_g.mean())
# volume IoU on shared grid
iou = float("nan")
if with_vol: # optional (slow; meaningless for non-watertight GT, see memory uv-flip note)
gp_vox = _voxelize_solid(pred_mesh)
gg_vox = _voxelize_solid(gt_mesh)
inter = np.logical_and(gp_vox, gg_vox).sum()
union = np.logical_or(gp_vox, gg_vox).sum()
iou = float(inter) / float(union) if union > 0 else 0.0
# ADDED: extent-relative metrics so FB150 (GT max extent 0.9) is comparable with
# Toys/Omni (1.0): distances divided by the GT max bbox extent (tau scaled by it).
ext = float((np.asarray(gt_mesh.vertices).max(0) - np.asarray(gt_mesh.vertices).min(0)).max())
rel = {}
for tau in TAUS:
t = tau * ext
pr_, rc_ = (d_p2g < t).float().mean(), (d_g2p < t).float().mean()
rel[tau] = float(2 * pr_ * rc_ / (pr_ + rc_ + 1e-12))
return {
"gt_extent": ext,
"cd_l1_rel": float(cd_l1) / ext,
"f01_rel": rel[0.01], "f02_rel": rel[0.02], "f05_rel": rel[0.05],
"cd_l1": float(cd_l1),
"cd_l2": float(cd_l2),
"fscore": {str(t): fscore[t] for t in TAUS},
"f01": fscore[0.01]["f"],
"f02": fscore[0.02]["f"], # headline
"f05": fscore[0.05]["f"],
"normal_consistency": float(nc),
"vol_iou": iou,
}
|