use core::panic;
use kd_tree::{self, KdPoint, KdTree};
use num_traits::{self};
use pasture_core::containers::{
BorrowedBuffer, BorrowedBufferExt, BorrowedMutBufferExt, HashMapBuffer, OwningBuffer,
};
use pasture_core::layout::{attributes::POSITION_3D, PointType};
use pasture_core::nalgebra::{DMatrix, Vector3};
use std::result::Result;
pub fn compute_normals<'a, T: BorrowedBuffer<'a>, P: PointType + KdPoint + Copy>(
point_cloud: &'a T,
k_nn: usize,
) -> Vec<(Vector3<f64>, f64)>
where
P::Scalar: num_traits::Float,
{
if point_cloud.len() < 3 {
panic!("The point cloud is too small. Please use a point cloud that has 3 or more points!");
}
if k_nn < 3 {
panic!("The k nearest neigbors attribute is too small!");
}
let mut points_with_normals_curvature = vec![];
let mut points: Vec<[f64; 3]> = vec![];
for point in point_cloud.view_attribute::<Vector3<f64>>(&POSITION_3D) {
points.push(*point.as_ref());
}
let cloud_as_kd_tree = KdTree::build_by_ordered_float(points);
for point in point_cloud.view_attribute::<Vector3<f64>>(&POSITION_3D) {
let r: &[f64; 3] = point.as_ref();
let nearest_points = cloud_as_kd_tree.nearests(r, k_nn);
let mut k_nn_buffer =
HashMapBuffer::with_capacity(nearest_points.len(), point_cloud.point_layout().clone());
k_nn_buffer.resize(nearest_points.len());
for (index, point) in nearest_points.iter().enumerate() {
k_nn_buffer.view_attribute_mut(&POSITION_3D).set_at(
index,
Vector3::new(point.item[0], point.item[1], point.item[2]),
);
}
let (surface_normal, curvature) = normal_estimation(&k_nn_buffer);
points_with_normals_curvature.push((surface_normal, curvature));
}
points_with_normals_curvature
}
fn is_dense<'a, T: BorrowedBuffer<'a>>(point_cloud: &'a T) -> bool {
for point in point_cloud.view_attribute::<Vector3<f64>>(&POSITION_3D) {
if point.x.is_nan() || point.y.is_nan() || point.z.is_nan() {
return false;
}
}
true
}
fn is_finite(point: &Vector3<f64>) -> bool {
if point.x.is_finite() && point.y.is_finite() && point.z.is_finite() {
return true;
}
false
}
pub fn compute_centroid<'a, T: BorrowedBuffer<'a>>(point_cloud: &'a T) -> Vector3<f64> {
if point_cloud.is_empty() {
panic!("The point cloud is empty!");
}
let mut centroid = Vector3::<f64>::zeros();
let mut temp_centroid = Vector3::<f64>::zeros();
if is_dense(point_cloud) {
for point in point_cloud.view_attribute::<Vector3<f64>>(&POSITION_3D) {
temp_centroid[0] += point.x;
temp_centroid[1] += point.y;
temp_centroid[2] += point.z;
}
centroid[0] = temp_centroid[0] / point_cloud.len() as f64;
centroid[1] = temp_centroid[1] / point_cloud.len() as f64;
centroid[2] = temp_centroid[2] / point_cloud.len() as f64;
} else {
let mut points_in_cloud = 0;
for point in point_cloud.view_attribute::<Vector3<f64>>(&POSITION_3D) {
if is_finite(&point) {
temp_centroid[0] += point.x;
temp_centroid[1] += point.y;
temp_centroid[2] += point.z;
points_in_cloud += 1;
}
}
centroid[0] = temp_centroid[0] / points_in_cloud as f64;
centroid[1] = temp_centroid[1] / points_in_cloud as f64;
centroid[2] = temp_centroid[2] / points_in_cloud as f64;
}
centroid
}
fn compute_covariance_matrix<'a, T: BorrowedBuffer<'a>>(
point_cloud: &'a T,
) -> Result<DMatrix<f64>, &'static str> {
let mut covariance_matrix = DMatrix::<f64>::zeros(3, 3);
let mut point_count = 0;
let centroid = compute_centroid(point_cloud);
if is_dense(point_cloud) {
point_count = point_cloud.len();
let mut diff_mean = Vector3::<f64>::zeros();
for point in point_cloud.view_attribute::<Vector3<f64>>(&POSITION_3D) {
diff_mean[0] = point.x - centroid[0];
diff_mean[1] = point.y - centroid[1];
diff_mean[2] = point.z - centroid[2];
covariance_matrix[(1, 1)] += diff_mean[1] * diff_mean[1];
covariance_matrix[(1, 2)] += diff_mean[1] * diff_mean[2];
covariance_matrix[(2, 2)] += diff_mean[2] * diff_mean[2];
let diff_x = diff_mean[0];
diff_mean.iter_mut().for_each(|x| *x *= diff_x);
covariance_matrix[(0, 0)] += diff_mean[0];
covariance_matrix[(0, 1)] += diff_mean[1];
covariance_matrix[(0, 2)] += diff_mean[2];
}
} else {
for point in point_cloud.view_attribute::<Vector3<f64>>(&POSITION_3D) {
if !is_finite(&point) {
continue;
}
let mut diff_mean = Vector3::<f64>::zeros();
diff_mean[0] = point.x - centroid[0];
diff_mean[1] = point.y - centroid[1];
diff_mean[2] = point.z - centroid[2];
covariance_matrix[(1, 1)] += diff_mean[1] * diff_mean[1];
covariance_matrix[(1, 2)] += diff_mean[1] * diff_mean[2];
covariance_matrix[(2, 2)] += diff_mean[2] * diff_mean[2];
let diff_x = diff_mean[0];
diff_mean.iter_mut().for_each(|x| *x *= diff_x);
covariance_matrix[(0, 0)] += diff_mean[0];
covariance_matrix[(0, 1)] += diff_mean[1];
covariance_matrix[(0, 2)] += diff_mean[2];
point_count += 1;
}
}
if point_count < 3 {
return Err("The number of valid (finite and non-NaN values) points in a k nearest neighborhood is not enough to span a plane!");
}
covariance_matrix[(1, 0)] = covariance_matrix[(0, 1)];
covariance_matrix[(2, 0)] = covariance_matrix[(0, 2)];
covariance_matrix[(2, 1)] = covariance_matrix[(1, 2)];
Ok(covariance_matrix)
}
fn solve_polynomial_quadratic(coefficient_2: f64, coefficient_1: f64) -> Vector3<f64> {
let mut eigen_values = Vector3::<f64>::zeros();
eigen_values[0] = 0.0;
let mut delta = coefficient_2 * coefficient_2 - 4.0 * coefficient_1;
if delta < 0.0 {
delta = 0.0;
}
let sqrt_delta = f64::sqrt(delta);
eigen_values[2] = 0.5 * (coefficient_2 + sqrt_delta);
eigen_values[1] = 0.5 * (coefficient_2 - sqrt_delta);
eigen_values
}
fn solve_polynomial(covariance_matrix: &DMatrix<f64>) -> Vector3<f64> {
let coefficient_0 = covariance_matrix[(0, 0)]
* covariance_matrix[(1, 1)]
* covariance_matrix[(2, 2)]
+ 2.0 * covariance_matrix[(0, 1)] * covariance_matrix[(0, 2)] * covariance_matrix[(1, 2)]
- covariance_matrix[(0, 0)] * covariance_matrix[(1, 2)] * covariance_matrix[(1, 2)]
- covariance_matrix[(1, 1)] * covariance_matrix[(0, 2)] * covariance_matrix[(0, 2)]
- covariance_matrix[(2, 2)] * covariance_matrix[(0, 1)] * covariance_matrix[(0, 1)];
let coefficient_1 = covariance_matrix[(0, 0)] * covariance_matrix[(1, 1)]
- covariance_matrix[(0, 1)] * covariance_matrix[(0, 1)]
+ covariance_matrix[(0, 0)] * covariance_matrix[(2, 2)]
- covariance_matrix[(0, 2)] * covariance_matrix[(0, 2)]
+ covariance_matrix[(1, 1)] * covariance_matrix[(2, 2)]
- covariance_matrix[(1, 2)] * covariance_matrix[(1, 2)];
let coefficient_2 =
covariance_matrix[(0, 0)] + covariance_matrix[(1, 1)] + covariance_matrix[(2, 2)];
if coefficient_0.abs() < std::f64::EPSILON {
solve_polynomial_quadratic(coefficient_2, coefficient_1)
} else {
let mut eigen_values = Vector3::<f64>::zeros();
let one_third = 1.0 / 3.0;
let sqrt_3 = f64::sqrt(3.0);
let coefficient_2_third = coefficient_2 * one_third;
let mut alpha_third = (coefficient_1 - coefficient_2 * coefficient_2_third) * one_third;
if alpha_third > 0.0 {
alpha_third = 0.0;
}
let half_beta = 0.5
* (coefficient_0
+ coefficient_2_third
* (2.0 * coefficient_2_third * coefficient_2_third - coefficient_1));
let mut q = half_beta * half_beta + alpha_third * alpha_third * alpha_third;
if q > 0.0 {
q = 0.0;
}
let rho = f64::sqrt(-alpha_third);
let theta = f64::atan2(f64::sqrt(-q), half_beta) * one_third;
let cosine_of_theta = f64::cos(theta);
let sine_of_theta = f64::sin(theta);
eigen_values[0] = coefficient_2_third + 2.0 * rho * cosine_of_theta;
eigen_values[1] = coefficient_2_third - rho * (cosine_of_theta + sqrt_3 * sine_of_theta);
eigen_values[2] = coefficient_2_third - rho * (cosine_of_theta - sqrt_3 * sine_of_theta);
eigen_values
.as_mut_slice()
.sort_by(|a, b| a.partial_cmp(b).unwrap());
if eigen_values[0] <= 0.0 {
solve_polynomial_quadratic(coefficient_2, coefficient_1)
} else {
eigen_values
}
}
}
fn get_largest_eigen_vector(scaled_matrix: &DMatrix<f64>) -> Vector3<f64> {
let rows = vec![
scaled_matrix.row(0).cross(&scaled_matrix.row(1)),
scaled_matrix.row(0).cross(&scaled_matrix.row(2)),
scaled_matrix.row(1).cross(&scaled_matrix.row(2)),
];
let mut cross_product = DMatrix::<f64>::zeros(3, 3);
for it in cross_product.row_iter_mut().zip(rows) {
let (mut cross_row, row) = it;
cross_row.copy_from(&row);
}
let mut largest_eigen_vec = cross_product.row(0);
for row in cross_product.row_iter() {
if row.norm() > largest_eigen_vec.norm() {
largest_eigen_vec = row;
}
}
let mut eigen_vector = Vector3::<f64>::zeros();
for i in 0..eigen_vector.len() {
eigen_vector[i] = largest_eigen_vec[i];
}
eigen_vector
}
fn eigen_3x3(covariance_matrix: &DMatrix<f64>) -> (f64, Vector3<f64>) {
let scale = covariance_matrix.abs().max();
let mut covariance_matrix_scaled = DMatrix::<f64>::zeros(3, 3);
for i in 0..covariance_matrix.len() {
covariance_matrix_scaled[i] = covariance_matrix[i] / scale;
}
for (index, value) in covariance_matrix.iter().enumerate() {
covariance_matrix_scaled[index] = value / scale;
}
let eigen_values = solve_polynomial(covariance_matrix);
let eigen_value = eigen_values[0] * scale;
covariance_matrix_scaled
.diagonal()
.iter_mut()
.for_each(|x| *x -= eigen_values[0]);
let eigen_vector = get_largest_eigen_vector(&covariance_matrix_scaled);
(eigen_value, eigen_vector)
}
fn solve_plane_parameter(covariance_matrix: &DMatrix<f64>) -> (Vector3<f64>, f64) {
let (eigen_value, eigen_vector) = eigen_3x3(covariance_matrix);
let eigen_sum = covariance_matrix[0] + covariance_matrix[4] + covariance_matrix[8];
let curvature = if eigen_sum != 0.0 {
(eigen_value / eigen_sum).abs()
} else {
0.0
};
(eigen_vector, curvature)
}
fn normal_estimation<'a, T: BorrowedBuffer<'a>>(point_cloud: &'a T) -> (Vector3<f64>, f64) {
let covariance_matrix = compute_covariance_matrix(point_cloud).unwrap();
let (eigen_vector, curvature) = solve_plane_parameter(&covariance_matrix);
(eigen_vector, curvature)
}
#[cfg(test)]
mod tests {
use pasture_core::{containers::VectorBuffer, nalgebra::Matrix3, nalgebra::Vector3};
use pasture_derive::PointType;
use super::*;
#[repr(C, packed)]
#[derive(PointType, Debug, Clone, Copy, bytemuck::AnyBitPattern, bytemuck::NoUninit)]
pub struct SimplePoint {
#[pasture(BUILTIN_POSITION_3D)]
pub position: Vector3<f64>,
#[pasture(BUILTIN_INTENSITY)]
pub intensity: u16,
}
impl KdPoint for SimplePoint {
type Scalar = f64;
type Dim = typenum::U3;
fn at(&self, k: usize) -> f64 {
let position = self.position;
position[k]
}
}
#[test]
fn test_compute_normal_sub() {
let points: Vec<SimplePoint> = vec![
SimplePoint {
position: Vector3::new(1.0, 0.0, 0.0),
intensity: 42,
},
SimplePoint {
position: Vector3::new(0.0, 1.0, 0.0),
intensity: 84,
},
SimplePoint {
position: Vector3::new(1.0, 1.0, 0.0),
intensity: 84,
},
SimplePoint {
position: Vector3::new(-1.0, 0.0, 0.0),
intensity: 84,
},
];
let interleaved = points.into_iter().collect::<VectorBuffer>();
let centroid = compute_centroid(&interleaved);
let result_centroid = Vector3::<f64>::new(0.25, 0.5, 0.0);
assert_eq!(centroid, result_centroid);
let covariance_matrix = compute_covariance_matrix(&interleaved).unwrap();
let result = Matrix3::new(
0.6875 * 4.0,
0.125 * 4.0,
0.0,
0.125 * 4.0,
0.25 * 4.0,
0.0,
0.0,
0.0,
0.0,
);
assert_eq!(covariance_matrix, result);
let (normal_vector, curvature) = solve_plane_parameter(&covariance_matrix);
assert_eq!(normal_vector[0], 0.0);
assert_eq!(normal_vector[1], 0.0);
assert_ne!(normal_vector[2], 0.0);
assert_eq!(curvature, 0.0);
}
#[test]
fn test_covariance_error() {
let points: Vec<SimplePoint> = vec![
SimplePoint {
position: Vector3::new(f64::NAN, 0.0, 0.0),
intensity: 42,
},
SimplePoint {
position: Vector3::new(0.0, 1.0, f64::NAN),
intensity: 84,
},
SimplePoint {
position: Vector3::new(1.0, 1.0, f64::NAN),
intensity: 84,
},
SimplePoint {
position: Vector3::new(-1.0, f64::NAN, 0.0),
intensity: 84,
},
];
let interleaved = points.into_iter().collect::<VectorBuffer>();
let result = compute_covariance_matrix(&interleaved);
let expected_result = Err("The number of valid (finite and non-NaN values) points in a k nearest neighborhood is not enough to span a plane!");
assert_eq!(result, expected_result);
}
#[test]
fn test_compute_normal() {
let points: Vec<SimplePoint> = vec![
SimplePoint {
position: Vector3::new(1.0, 0.0, 0.0),
intensity: 42,
},
SimplePoint {
position: Vector3::new(0.0, 1.0, 0.0),
intensity: 84,
},
SimplePoint {
position: Vector3::new(1.0, 1.0, 0.0),
intensity: 84,
},
SimplePoint {
position: Vector3::new(-1.0, 0.0, 0.0),
intensity: 84,
},
];
let interleaved = points.into_iter().collect::<VectorBuffer>();
let solution_vec = compute_normals::<VectorBuffer, SimplePoint>(&interleaved, 3);
for solution in solution_vec {
assert_eq!(solution.0[0], 0.0);
assert_eq!(solution.0[1], 0.0);
assert_ne!(solution.0[2], 0.0);
assert_eq!(solution.1, 0.0);
}
}
#[test]
#[should_panic(
expected = "The point cloud is too small. Please use a point cloud that has 3 or more points!"
)]
fn test_compute_normal_not_enough_points_1() {
let points: Vec<SimplePoint> = vec![SimplePoint {
position: Vector3::new(1.0, 0.0, 0.0),
intensity: 42,
}];
let interleaved = points.into_iter().collect::<VectorBuffer>();
let _solution_vec = compute_normals::<VectorBuffer, SimplePoint>(&interleaved, 3);
}
#[test]
#[should_panic(
expected = "The point cloud is too small. Please use a point cloud that has 3 or more points!"
)]
fn test_compute_normal_not_enough_points_2() {
let points: Vec<SimplePoint> = vec![
SimplePoint {
position: Vector3::new(1.0, 0.0, 0.0),
intensity: 42,
},
SimplePoint {
position: Vector3::new(0.0, 1.0, 0.0),
intensity: 84,
},
];
let interleaved = points.into_iter().collect::<VectorBuffer>();
let _solution_vec = compute_normals::<VectorBuffer, SimplePoint>(&interleaved, 3);
}
#[test]
#[should_panic(expected = "The k nearest neigbors attribute is too small!")]
fn test_compute_normal_knn_too_small_1() {
let points: Vec<SimplePoint> = vec![
SimplePoint {
position: Vector3::new(1.0, 0.0, 0.0),
intensity: 42,
},
SimplePoint {
position: Vector3::new(0.0, 1.0, 0.0),
intensity: 84,
},
SimplePoint {
position: Vector3::new(1.0, 1.0, 0.0),
intensity: 84,
},
SimplePoint {
position: Vector3::new(-1.0, 0.0, 0.0),
intensity: 84,
},
];
let interleaved = points.into_iter().collect::<VectorBuffer>();
let _solution_vec = compute_normals::<VectorBuffer, SimplePoint>(&interleaved, 1);
}
#[test]
#[should_panic(expected = "The k nearest neigbors attribute is too small!")]
fn test_compute_normal_knn_too_small_2() {
let points: Vec<SimplePoint> = vec![
SimplePoint {
position: Vector3::new(1.0, 0.0, 0.0),
intensity: 42,
},
SimplePoint {
position: Vector3::new(0.0, 1.0, 0.0),
intensity: 84,
},
SimplePoint {
position: Vector3::new(1.0, 1.0, 0.0),
intensity: 84,
},
SimplePoint {
position: Vector3::new(-1.0, 0.0, 0.0),
intensity: 84,
},
];
let interleaved = points.into_iter().collect::<VectorBuffer>();
let _solution_vec = compute_normals::<VectorBuffer, SimplePoint>(&interleaved, 2);
}
}