40 lines
1.6 KiB
Python
40 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
import copy
|
|
import numpy as np
|
|
import open3d as o3d
|
|
|
|
|
|
class CurvatureEstimator:
|
|
"""Estimate local point-cloud curvature using neighborhood covariance."""
|
|
|
|
def compute(self, pointcloud: o3d.geometry.PointCloud, radius: float = 2.0, max_nn: int = 30) -> np.ndarray:
|
|
pointcloud.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=max_nn))
|
|
points = np.asarray(pointcloud.points)
|
|
kdtree = o3d.geometry.KDTreeFlann(pointcloud)
|
|
curvatures = np.zeros(len(points), dtype=np.float64)
|
|
|
|
for i, point in enumerate(points):
|
|
_, indices, _ = kdtree.search_hybrid_vector_3d(point, radius, max_nn)
|
|
if len(indices) < 3:
|
|
curvatures[i] = 1.0
|
|
continue
|
|
|
|
neighbors = points[indices]
|
|
cov = np.cov((neighbors - neighbors.mean(axis=0)).T)
|
|
eigenvalues = np.linalg.eigvalsh(cov)
|
|
denom = eigenvalues.sum()
|
|
curvatures[i] = float(eigenvalues[0] / denom) if denom > 0 else 0.0
|
|
|
|
return curvatures
|
|
|
|
def colorize(self, pointcloud: o3d.geometry.PointCloud, curvatures: np.ndarray, color_scale: float = 3.0) -> o3d.geometry.PointCloud:
|
|
pcd = copy.deepcopy(pointcloud)
|
|
normalized = curvatures / (np.max(curvatures) + 1e-12)
|
|
normalized = np.clip(normalized * color_scale, 0, 1)
|
|
colors = np.zeros((len(normalized), 3))
|
|
colors[:, 0] = normalized
|
|
colors[:, 1] = 1 - normalized
|
|
pcd.colors = o3d.utility.Vector3dVector(colors)
|
|
return pcd
|