Spaces:
Running on Zero
Running on Zero
File size: 3,740 Bytes
e340a84 | 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 | import numpy as np
from scipy.spatial import cKDTree
def similarity_align(src, dst, with_scale=True):
src = np.asarray(src, dtype=np.float64)
dst = np.asarray(dst, dtype=np.float64)
if src.shape != dst.shape or src.ndim != 2 or src.shape[1] != 3:
raise ValueError("Expected Nx3 source and target point sets")
if len(src) < 3:
return 1.0, np.eye(3), np.zeros(3)
src_mean = src.mean(axis=0)
dst_mean = dst.mean(axis=0)
src_centered = src - src_mean
dst_centered = dst - dst_mean
cov = (dst_centered.T @ src_centered) / len(src)
U, D, Vt = np.linalg.svd(cov)
S = np.eye(3)
if np.linalg.det(U @ Vt) < 0:
S[-1, -1] = -1.0
R = U @ S @ Vt
if with_scale:
var = np.mean(np.sum(src_centered ** 2, axis=1))
scale = float(np.trace(np.diag(D) @ S) / max(var, 1e-12))
else:
scale = 1.0
t = dst_mean - scale * (R @ src_mean)
return scale, R, t
def transform_points(points, scale, R, t):
return (scale * (R @ points.T)).T + t[None]
def ate_rmse(pred_xyz, gt_xyz, align_scale=True):
scale, R, t = similarity_align(pred_xyz, gt_xyz, with_scale=align_scale)
pred_aligned = transform_points(pred_xyz, scale, R, t)
err = np.linalg.norm(pred_aligned - gt_xyz, axis=1)
return {
"ate_rmse": float(np.sqrt(np.mean(err ** 2))),
"ate_mean": float(np.mean(err)),
"ate_median": float(np.median(err)),
"num_pose_pairs": int(len(err)),
"align_scale": bool(align_scale),
"sim3_scale": float(scale),
"sim3_rotation": R.tolist(),
"sim3_translation": t.tolist(),
}
def _voxel_downsample(points, voxel_size):
if voxel_size is None:
return points
voxel_size = float(voxel_size)
if voxel_size <= 0 or len(points) == 0:
return points
coords = np.floor(points / voxel_size).astype(np.int64)
_, keep = np.unique(coords, axis=0, return_index=True)
keep.sort()
return points[keep]
def _sample_points(points, max_points, seed):
if max_points is None or len(points) <= int(max_points):
return points
rng = np.random.default_rng(seed)
keep = rng.choice(len(points), size=int(max_points), replace=False)
return points[keep]
def prepare_pointcloud(points, max_points=None, voxel_size=None, seed=0):
points = np.asarray(points, dtype=np.float64).reshape(-1, 3)
if len(points) == 0:
return points
valid = np.isfinite(points).all(axis=1)
points = points[valid]
points = _voxel_downsample(points, voxel_size)
points = _sample_points(points, max_points, seed)
return points
def chamfer_and_f1(
pred_points, gt_points, threshold=0.25, max_points=None, voxel_size=None, seed=0
):
pred = prepare_pointcloud(
pred_points, max_points=max_points, voxel_size=voxel_size, seed=seed
)
gt = prepare_pointcloud(
gt_points, max_points=max_points, voxel_size=voxel_size, seed=seed + 1
)
if len(pred) == 0 or len(gt) == 0:
return None
pred_tree = cKDTree(pred)
gt_tree = cKDTree(gt)
dist_pred_to_gt, _ = gt_tree.query(pred, k=1)
dist_gt_to_pred, _ = pred_tree.query(gt, k=1)
acc = float(np.mean(dist_pred_to_gt))
comp = float(np.mean(dist_gt_to_pred))
precision = float(np.mean(dist_pred_to_gt < threshold))
recall = float(np.mean(dist_gt_to_pred < threshold))
denom = precision + recall
f1 = 0.0 if denom <= 0 else float(2.0 * precision * recall / denom)
return {
"cd": float(acc + comp),
"acc": acc,
"comp": comp,
"f1": f1,
"f1_threshold": float(threshold),
"num_pred_points": int(len(pred)),
"num_gt_points": int(len(gt)),
}
|