212 lines
8.5 KiB
Python
212 lines
8.5 KiB
Python
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
import open3d as o3d
|
||
import torch
|
||
import trimesh
|
||
from scipy.spatial import KDTree
|
||
|
||
|
||
@dataclass
|
||
class DistanceStats:
|
||
mean: float
|
||
max: float
|
||
rmse: float
|
||
min: float
|
||
std: float
|
||
|
||
|
||
class Evaluator:
|
||
"""Evaluation metrics for point-cloud workflows."""
|
||
|
||
def chamfer_distance(
|
||
self,
|
||
pointcloud_a: o3d.geometry.PointCloud,
|
||
pointcloud_b: o3d.geometry.PointCloud,
|
||
sqrt: bool = True,
|
||
) -> tuple[float, float, float]:
|
||
points_a = np.asarray(pointcloud_a.points)
|
||
points_b = np.asarray(pointcloud_b.points)
|
||
if len(points_a) == 0 or len(points_b) == 0:
|
||
return float("inf"), 0.0, 0.0
|
||
|
||
tree_a = KDTree(points_a)
|
||
tree_b = KDTree(points_b)
|
||
dist_a, _ = tree_b.query(points_a, k=1)
|
||
dist_b, _ = tree_a.query(points_b, k=1)
|
||
cd_a = float(np.mean(dist_a ** 2))
|
||
cd_b = float(np.mean(dist_b ** 2))
|
||
cd = cd_a + cd_b
|
||
return (float(np.sqrt(cd)), cd_a, cd_b) if sqrt else (float(cd), cd_a, cd_b)
|
||
|
||
def point_to_mesh_distance_trimesh(self, points: np.ndarray, mesh_path: str) -> tuple[np.ndarray, DistanceStats]:
|
||
mesh = trimesh.load(mesh_path, force="mesh")
|
||
_, distances, _ = mesh.nearest.on_surface(points)
|
||
distances = np.asarray(distances, dtype=np.float64)
|
||
return distances, self.distance_stats(distances)
|
||
|
||
def point_to_mesh_distance_torch(
|
||
self,
|
||
points: np.ndarray,
|
||
mesh: trimesh.Trimesh,
|
||
chunk_size: int = 2048,
|
||
triangle_chunk_size: int | None = 4096,
|
||
) -> tuple[np.ndarray, DistanceStats]:
|
||
vertices, faces = self._mesh_arrays(mesh)
|
||
return self.point_to_triangles_distance_torch(
|
||
points,
|
||
vertices,
|
||
faces,
|
||
point_chunk_size=chunk_size,
|
||
triangle_chunk_size=triangle_chunk_size,
|
||
)
|
||
|
||
def point_to_triangles_distance_torch(
|
||
self,
|
||
points: np.ndarray,
|
||
vertices: np.ndarray,
|
||
faces: np.ndarray,
|
||
point_chunk_size: int = 2048,
|
||
triangle_chunk_size: int | None = 4096,
|
||
) -> tuple[np.ndarray, DistanceStats]:
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
dtype = torch.float64
|
||
|
||
# vertices/triangles 上 GPU（mesh 通常几百MB，可以接受）
|
||
vertices_t = torch.as_tensor(np.asarray(vertices, dtype=np.float64), dtype=dtype, device=device)
|
||
faces_t = torch.as_tensor(np.asarray(faces, dtype=np.int64), dtype=torch.long, device=device)
|
||
triangles = vertices_t[faces_t]
|
||
del vertices_t, faces_t # 释放原始 vertices/faces，只保留 triangles
|
||
|
||
# points 保留在 CPU，按 point_chunk 逐批传上 GPU，避免 10 亿点全量上 GPU OOM
|
||
points_np = np.asarray(points, dtype=np.float64)
|
||
n_points = points_np.shape[0]
|
||
|
||
distances = []
|
||
with torch.no_grad():
|
||
for start in range(0, n_points, point_chunk_size):
|
||
chunk = torch.as_tensor(
|
||
points_np[start:start + point_chunk_size], dtype=dtype, device=device
|
||
)
|
||
if triangle_chunk_size is None or triangle_chunk_size <= 0:
|
||
dist2 = self._point_to_triangle_distance_squared(chunk, triangles)
|
||
min_dist2 = dist2.min(dim=1).values
|
||
else:
|
||
min_dist2 = torch.full(
|
||
(chunk.shape[0],),
|
||
torch.inf,
|
||
dtype=dtype,
|
||
device=device,
|
||
)
|
||
for triangle_start in range(0, triangles.shape[0], triangle_chunk_size):
|
||
triangle_chunk = triangles[triangle_start:triangle_start + triangle_chunk_size]
|
||
dist2 = self._point_to_triangle_distance_squared(chunk, triangle_chunk)
|
||
min_dist2 = torch.minimum(min_dist2, dist2.min(dim=1).values)
|
||
|
||
distances.append(torch.sqrt(min_dist2).cpu().numpy())
|
||
|
||
all_distances = np.concatenate(distances) if distances else np.array([], dtype=np.float64)
|
||
return all_distances, self.distance_stats(all_distances)
|
||
|
||
def colorize_distances(self, points: np.ndarray, distances: np.ndarray) -> o3d.geometry.PointCloud:
|
||
if len(points) != len(distances):
|
||
raise ValueError("points and distances must have the same length.")
|
||
pointcloud = o3d.geometry.PointCloud()
|
||
pointcloud.points = o3d.utility.Vector3dVector(np.asarray(points, dtype=np.float64))
|
||
|
||
if len(distances) == 0:
|
||
pointcloud.colors = o3d.utility.Vector3dVector(np.empty((0, 3), dtype=np.float64))
|
||
return pointcloud
|
||
|
||
lower = float(np.min(distances))
|
||
upper = float(np.max(distances))
|
||
normalized = (distances - lower) / (upper - lower + 1e-12)
|
||
colors = self._blue_to_red_colormap(normalized)
|
||
pointcloud.colors = o3d.utility.Vector3dVector(colors)
|
||
return pointcloud
|
||
|
||
def distance_stats(self, distances: np.ndarray) -> DistanceStats:
|
||
if len(distances) == 0:
|
||
return DistanceStats(0.0, 0.0, 0.0, 0.0, 0.0)
|
||
return DistanceStats(
|
||
mean=float(np.mean(distances)),
|
||
max=float(np.max(distances)),
|
||
rmse=float(np.sqrt(np.mean(distances ** 2))),
|
||
min=float(np.min(distances)),
|
||
std=float(np.std(distances)),
|
||
)
|
||
|
||
def registration_summary(self, transformation: np.ndarray, fitness: float, rmse: float) -> dict[str, Any]:
|
||
return {
|
||
"fitness": fitness,
|
||
"rmse": rmse,
|
||
"transformation": transformation,
|
||
}
|
||
|
||
def _mesh_arrays(self, mesh: trimesh.Trimesh) -> tuple[np.ndarray, np.ndarray]:
|
||
vertices = np.asarray(mesh.vertices, dtype=np.float64)
|
||
faces = np.asarray(mesh.faces, dtype=np.int64)
|
||
if len(vertices) == 0 or len(faces) == 0:
|
||
raise ValueError("Mesh must contain vertices and faces.")
|
||
return vertices, faces
|
||
|
||
def _blue_to_red_colormap(self, values: np.ndarray) -> np.ndarray:
|
||
values = np.clip(np.asarray(values, dtype=np.float64), 0.0, 1.0)
|
||
colors = np.zeros((len(values), 3), dtype=np.float64)
|
||
colors[:, 0] = np.clip(1.8 * values - 0.2, 0.0, 1.0)
|
||
colors[:, 1] = np.clip(1.0 - np.abs(values - 0.5) * 2.0, 0.0, 1.0)
|
||
colors[:, 2] = np.clip(1.2 - 1.8 * values, 0.0, 1.0)
|
||
return colors
|
||
|
||
def _point_to_triangle_distance_squared(self, points: torch.Tensor, triangles: torch.Tensor) -> torch.Tensor:
|
||
p = points[:, None, :]
|
||
a = triangles[None, :, 0, :]
|
||
b = triangles[None, :, 1, :]
|
||
c = triangles[None, :, 2, :]
|
||
|
||
ab = b - a
|
||
ac = c - a
|
||
ap = p - a
|
||
|
||
d1 = (ab * ap).sum(dim=-1)
|
||
d2 = (ac * ap).sum(dim=-1)
|
||
inf = torch.full_like(d1, torch.inf)
|
||
dist2 = torch.where((d1 <= 0.0) & (d2 <= 0.0), (ap * ap).sum(dim=-1), inf)
|
||
|
||
bp = p - b
|
||
d3 = (ab * bp).sum(dim=-1)
|
||
d4 = (ac * bp).sum(dim=-1)
|
||
dist2 = torch.minimum(dist2, torch.where((d3 >= 0.0) & (d4 <= d3), (bp * bp).sum(dim=-1), inf))
|
||
|
||
vc = d1 * d4 - d3 * d2
|
||
mask = (vc <= 0.0) & (d1 >= 0.0) & (d3 <= 0.0)
|
||
v = d1 / (d1 - d3).clamp_min(torch.finfo(points.dtype).eps)
|
||
closest = a + v[..., None] * ab
|
||
dist2 = torch.minimum(dist2, torch.where(mask, ((p - closest) ** 2).sum(dim=-1), inf))
|
||
|
||
cp = p - c
|
||
d5 = (ab * cp).sum(dim=-1)
|
||
d6 = (ac * cp).sum(dim=-1)
|
||
dist2 = torch.minimum(dist2, torch.where((d6 >= 0.0) & (d5 <= d6), (cp * cp).sum(dim=-1), inf))
|
||
|
||
vb = d5 * d2 - d1 * d6
|
||
mask = (vb <= 0.0) & (d2 >= 0.0) & (d6 <= 0.0)
|
||
w = d2 / (d2 - d6).clamp_min(torch.finfo(points.dtype).eps)
|
||
closest = a + w[..., None] * ac
|
||
dist2 = torch.minimum(dist2, torch.where(mask, ((p - closest) ** 2).sum(dim=-1), inf))
|
||
|
||
va = d3 * d6 - d5 * d4
|
||
mask = (va <= 0.0) & ((d4 - d3) >= 0.0) & ((d5 - d6) >= 0.0)
|
||
w = (d4 - d3) / ((d4 - d3) + (d5 - d6)).clamp_min(torch.finfo(points.dtype).eps)
|
||
closest = b + w[..., None] * (c - b)
|
||
dist2 = torch.minimum(dist2, torch.where(mask, ((p - closest) ** 2).sum(dim=-1), inf))
|
||
|
||
denom = (va + vb + vc).clamp_min(torch.finfo(points.dtype).eps)
|
||
v = vb / denom
|
||
w = vc / denom
|
||
closest = a + ab * v[..., None] + ac * w[..., None]
|
||
return torch.minimum(dist2, ((p - closest) ** 2).sum(dim=-1))
|