239 lines
8.6 KiB
Python
239 lines
8.6 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
from scipy.spatial import cKDTree
|
|
|
|
from .roi import apply_passthrough, normalize_aabb
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class PreprocessResult:
|
|
raw_points: np.ndarray
|
|
filtered_points: np.ndarray
|
|
work_points: np.ndarray
|
|
work_normals: np.ndarray
|
|
work_curvature: np.ndarray
|
|
spacing: float
|
|
stats: dict
|
|
|
|
|
|
def estimate_spacing(points: np.ndarray, sample_size: int = 4000) -> float:
|
|
if len(points) < 2:
|
|
return 1.0
|
|
if len(points) > sample_size:
|
|
indices = np.linspace(0, len(points) - 1, sample_size, dtype=int)
|
|
sample = points[indices]
|
|
else:
|
|
sample = points
|
|
tree = cKDTree(sample)
|
|
distances, _ = tree.query(sample, k=min(2, len(sample)))
|
|
if distances.ndim == 1:
|
|
return float(np.median(distances))
|
|
return float(np.quantile(distances[:, -1], 0.3))
|
|
|
|
|
|
def statistical_outlier_filter(
|
|
points: np.ndarray,
|
|
k: int,
|
|
std_ratio: float,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
if len(points) <= k + 1:
|
|
mask = np.ones(len(points), dtype=bool)
|
|
return points, mask
|
|
tree = cKDTree(points)
|
|
distances, _ = tree.query(points, k=k + 1)
|
|
mean_dist = distances[:, 1:].mean(axis=1)
|
|
threshold = mean_dist.mean() + std_ratio * mean_dist.std(ddof=0)
|
|
mask = mean_dist <= threshold
|
|
return points[mask], mask
|
|
|
|
|
|
def radius_filter(
|
|
points: np.ndarray,
|
|
radius: float,
|
|
min_neighbors: int,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
if len(points) == 0:
|
|
return points, np.zeros(0, dtype=bool)
|
|
tree = cKDTree(points)
|
|
counts = tree.query_ball_point(points, r=radius, return_length=True)
|
|
mask = counts >= min_neighbors
|
|
return points[mask], mask
|
|
|
|
|
|
def voxel_downsample(points: np.ndarray, voxel_size: float) -> np.ndarray:
|
|
if len(points) == 0 or voxel_size <= 0:
|
|
return points.copy()
|
|
mins = points.min(axis=0)
|
|
grid = np.floor((points - mins) / voxel_size).astype(np.int64)
|
|
unique, inverse = np.unique(grid, axis=0, return_inverse=True)
|
|
sums = np.zeros((len(unique), 3), dtype=np.float64)
|
|
counts = np.bincount(inverse)
|
|
np.add.at(sums, inverse, points)
|
|
return sums / counts[:, None]
|
|
|
|
|
|
def estimate_normals(
|
|
points: np.ndarray,
|
|
k: int,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
if len(points) == 0:
|
|
return np.zeros((0, 3)), np.zeros(0)
|
|
neighbor_count = min(k + 1, len(points))
|
|
tree = cKDTree(points)
|
|
_, indices = tree.query(points, k=neighbor_count)
|
|
normals = np.zeros((len(points), 3), dtype=np.float64)
|
|
curvature = np.zeros(len(points), dtype=np.float64)
|
|
for idx, neigh in enumerate(indices):
|
|
neighborhood = points[np.atleast_1d(neigh)]
|
|
centered = neighborhood - neighborhood.mean(axis=0, keepdims=True)
|
|
cov = centered.T @ centered / max(len(centered), 1)
|
|
eigenvalues, eigenvectors = np.linalg.eigh(cov)
|
|
order = np.argsort(eigenvalues)
|
|
eigenvalues = eigenvalues[order]
|
|
eigenvectors = eigenvectors[:, order]
|
|
normals[idx] = eigenvectors[:, 0]
|
|
curvature[idx] = float(eigenvalues[0] / max(eigenvalues.sum(), 1e-9))
|
|
return normals, curvature
|
|
|
|
|
|
def bilateral_filter_points(
|
|
points: np.ndarray,
|
|
normals: np.ndarray,
|
|
*,
|
|
k: int,
|
|
sigma_d_mm: float,
|
|
sigma_n_rad: float,
|
|
iterations: int = 1,
|
|
) -> np.ndarray:
|
|
"""Point-cloud bilateral filter that smooths along estimated normals.
|
|
|
|
Each point is shifted along its own normal by the weighted average of
|
|
the signed normal-direction offsets from its k-nearest neighbours,
|
|
using a Gaussian on Euclidean distance (``sigma_d_mm``) and a
|
|
Gaussian on the angular deviation between point normals
|
|
(``sigma_n_rad``). The tangential position is preserved, so edges
|
|
survive while sensor noise along the surface normal is suppressed.
|
|
"""
|
|
if len(points) == 0 or iterations <= 0:
|
|
return points.copy()
|
|
effective_k = int(min(max(k, 2), len(points)))
|
|
sigma_d = max(float(sigma_d_mm), 1e-6)
|
|
sigma_n = max(float(sigma_n_rad), 1e-6)
|
|
inv_2sd2 = 1.0 / (2.0 * sigma_d * sigma_d)
|
|
inv_2sn2 = 1.0 / (2.0 * sigma_n * sigma_n)
|
|
|
|
current = points.astype(np.float64, copy=True)
|
|
for _ in range(int(iterations)):
|
|
tree = cKDTree(current)
|
|
distances, indices = tree.query(current, k=effective_k)
|
|
if distances.ndim == 1: # k == 1 edge case
|
|
distances = distances[:, None]
|
|
indices = indices[:, None]
|
|
neighbor_points = current[indices] # (N, k, 3)
|
|
neighbor_normals = normals[indices] # (N, k, 3)
|
|
offsets = neighbor_points - current[:, None, :] # (N, k, 3)
|
|
own_normals = normals[:, None, :] # (N, 1, 3)
|
|
normal_offset = np.einsum("nkd,nkd->nk", offsets, own_normals)
|
|
spatial_w = np.exp(-(distances ** 2) * inv_2sd2)
|
|
dots = np.clip(np.einsum("nkd,nkd->nk", neighbor_normals, own_normals), -1.0, 1.0)
|
|
angle = np.arccos(np.abs(dots))
|
|
normal_w = np.exp(-(angle ** 2) * inv_2sn2)
|
|
weights = spatial_w * normal_w
|
|
weight_sum = weights.sum(axis=1)
|
|
valid = weight_sum > 1e-12
|
|
delta = np.zeros(len(current), dtype=np.float64)
|
|
delta[valid] = (
|
|
(weights[valid] * normal_offset[valid]).sum(axis=1) / weight_sum[valid]
|
|
)
|
|
current = current + delta[:, None] * normals
|
|
return current
|
|
|
|
|
|
def preprocess_point_cloud(points: np.ndarray, config: dict) -> PreprocessResult:
|
|
preprocess = config["preprocess"]
|
|
flags: list[str] = []
|
|
spacing_estimate = estimate_spacing(
|
|
points,
|
|
sample_size=preprocess["spacing_sample_size"],
|
|
)
|
|
|
|
passthrough_cfg = preprocess.get("passthrough") or {}
|
|
passthrough_aabb = passthrough_cfg.get("aabb") if passthrough_cfg.get("enabled") else None
|
|
raw_count = int(len(points))
|
|
if passthrough_aabb is not None:
|
|
aabb = normalize_aabb(passthrough_aabb)
|
|
keep_mask = apply_passthrough(points, aabb)
|
|
kept = points[keep_mask]
|
|
if len(kept) >= max(64, raw_count // 1000):
|
|
passthrough_points = kept
|
|
else:
|
|
passthrough_points = points
|
|
flags.append("passthrough_empty")
|
|
else:
|
|
passthrough_points = points
|
|
|
|
stat_points, stat_mask = statistical_outlier_filter(
|
|
passthrough_points,
|
|
k=preprocess["stat_k"],
|
|
std_ratio=preprocess["stat_std_ratio"],
|
|
)
|
|
radius_points, radius_mask = radius_filter(
|
|
stat_points,
|
|
radius=preprocess["radius_mm"],
|
|
min_neighbors=preprocess["radius_min_neighbors"],
|
|
)
|
|
|
|
bilateral_cfg = preprocess.get("bilateral") or {}
|
|
bilateral_count = len(radius_points)
|
|
if bilateral_cfg.get("enabled") and len(radius_points) >= 8:
|
|
k_requested = int(bilateral_cfg.get("k", 24))
|
|
k_effective = min(k_requested, len(radius_points))
|
|
if k_effective < k_requested:
|
|
flags.append("bilateral_degraded")
|
|
pre_normals, _ = estimate_normals(radius_points, k=min(k_effective, 32))
|
|
radius_points = bilateral_filter_points(
|
|
radius_points,
|
|
pre_normals,
|
|
k=k_effective,
|
|
sigma_d_mm=float(bilateral_cfg.get("sigma_d_mm", max(2.0 * spacing_estimate, 2.0))),
|
|
sigma_n_rad=float(
|
|
bilateral_cfg.get(
|
|
"sigma_n_rad",
|
|
np.deg2rad(float(bilateral_cfg.get("sigma_n_deg", 15.0))),
|
|
)
|
|
),
|
|
iterations=int(bilateral_cfg.get("iterations", 1)),
|
|
)
|
|
|
|
work_points = voxel_downsample(radius_points, preprocess["voxel_mm"])
|
|
work_normals, work_curvature = estimate_normals(
|
|
work_points,
|
|
k=preprocess["normal_k"],
|
|
)
|
|
stats = {
|
|
"raw_count": raw_count,
|
|
"after_passthrough": int(len(passthrough_points)),
|
|
"passthrough_applied": passthrough_aabb is not None,
|
|
"after_statistical_filter": int(len(stat_points)),
|
|
"after_radius_filter": int(len(radius_points)),
|
|
"after_bilateral": int(bilateral_count),
|
|
"bilateral_applied": bool(bilateral_cfg.get("enabled") and len(radius_points) >= 8),
|
|
"downsampled_count": int(len(work_points)),
|
|
"spacing_mm": float(spacing_estimate),
|
|
"statistical_keep_ratio": float(np.mean(stat_mask)) if len(stat_mask) else 1.0,
|
|
"radius_keep_ratio": float(np.mean(radius_mask)) if len(radius_mask) else 1.0,
|
|
"flags": flags,
|
|
}
|
|
return PreprocessResult(
|
|
raw_points=points,
|
|
filtered_points=radius_points,
|
|
work_points=work_points,
|
|
work_normals=work_normals,
|
|
work_curvature=work_curvature,
|
|
spacing=spacing_estimate,
|
|
stats=stats,
|
|
)
|