feat: initial commit - massive-cloud point cloud processing system
This commit is contained in:
203
alg/point_porj/evaluation.py
Normal file
203
alg/point_porj/evaluation.py
Normal file
@@ -0,0 +1,203 @@
|
||||
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))
|
||||
Reference in New Issue
Block a user