204 lines
8.1 KiB
Python
204 lines
8.1 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_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]
|
|
points_t = torch.as_tensor(points, dtype=dtype, device=device)
|
|
|
|
distances = []
|
|
with torch.no_grad():
|
|
for start in range(0, points_t.shape[0], point_chunk_size):
|
|
chunk = points_t[start:start + point_chunk_size]
|
|
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=points_t.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))
|