Skip to content

Commit

Permalink
Improve bvh pruning
Browse files Browse the repository at this point in the history
  • Loading branch information
@Azkellas
Azkellas committed Jan 8, 2025
1 parent beea42b commit 18f0bee
Show file tree
Hide file tree
Showing 4 changed files with 240 additions and 167 deletions.
64 changes: 30 additions & 34 deletions mesh_to_sdf/src/bvh_ext.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,19 @@ pub trait AabbExt {
/// The minimum distance is the distance to the closest point on the box,
/// 0 if the point is inside the box.
/// The maximum distance is the distance to the furthest point in the box.
fn get_min_max_distances<V: Point>(&self, point: &V) -> (f32, f32);
fn get_min_max_distances_2<V: Point>(&self, point: &V) -> (f32, f32);
}

pub trait BvhNodeExt<V: Point> {
fn get_dist_2(&self, point: &V) -> f32;
}

impl AabbExt for Aabb<f32, 3> {
/// Returns the minimum and maximum distances to the box.
/// The minimum distance is the distance to the closest point on the box,
/// 0 if the point is inside the box.
/// The maximum distance is the distance to the furthest point in the box.
fn get_min_max_distances<V: Point>(&self, point: &V) -> (f32, f32) {
fn get_min_max_distances_2<V: Point>(&self, point: &V) -> (f32, f32) {
// Convert point to nalgebra point
let n_point: nalgebra::OPoint<f32, nalgebra::Const<3>> =
nalgebra::Point3::new(point.x(), point.y(), point.z());
Expand All @@ -39,42 +43,39 @@ impl AabbExt for Aabb<f32, 3> {
let furthest = center + signum.component_mul(&half_size);
let max_dist = (n_point - furthest).norm();

(min_dist, max_dist)
(min_dist * min_dist, max_dist * max_dist)
}
}

pub trait BvhDistance<V: Point, Shape: Bounded<f32, 3>> {
/// Traverses the [`Bvh`].
/// Returns a subset of `shapes` which are candidates for being the closest to `point`.
///
fn nearest_candidates(&self, origin: &V, shapes: &[Shape]) -> Vec<usize>
fn nearest_candidates(&self, origin: &V, shapes: &[Shape]) -> (usize, f32)
where
Self: core::marker::Sized;
}

impl<V: Point, Shape: Bounded<f32, 3>> BvhDistance<V, Shape> for Bvh<f32, 3> {
impl<V: Point, Shape: Bounded<f32, 3> + BvhNodeExt<V>> BvhDistance<V, Shape> for Bvh<f32, 3> {
/// Traverses the [`Bvh`].
/// Returns a subset of `shapes` which are candidates for being the closest to `point`.
///
fn nearest_candidates(&self, origin: &V, shapes: &[Shape]) -> Vec<usize> {
let mut indices = Vec::new();
fn nearest_candidates(&self, origin: &V, shapes: &[Shape]) -> (usize, f32) {
let mut best_candidate = (0, f32::MAX);
let mut best_min_distance = f32::MAX;
let mut best_max_distance = f32::MAX;
BvhNode::nearest_candidates_recursive(
&self.nodes,
0,
origin,
shapes,
&mut indices,
&mut best_candidate,
&mut best_min_distance,
&mut best_max_distance,
);

indices
.into_iter()
.filter(|(_, node_min)| *node_min <= best_max_distance)
.map(|(i, _)| i)
.collect()
// Return the real distance and not the squared distance.
(best_candidate.0, best_candidate.1.sqrt())
}
}

Expand All @@ -87,14 +88,16 @@ pub trait BvhTraverseDistance<V: Point, Shape: Bounded<f32, 3>> {
node_index: usize,
origin: &V,
shapes: &[Shape],
indices: &mut Vec<(usize, f32)>,
best_candidate: &mut (usize, f32),
best_min_distance: &mut f32,
best_max_distance: &mut f32,
) where
Self: core::marker::Sized;
}

impl<V: Point, Shape: Bounded<f32, 3>> BvhTraverseDistance<V, Shape> for BvhNode<f32, 3> {
impl<V: Point, Shape: Bounded<f32, 3> + BvhNodeExt<V>> BvhTraverseDistance<V, Shape>
for BvhNode<f32, 3>
{
/// Traverses the [`Bvh`] recursively and returns all shapes whose [`Aabb`] countains
/// a candidate shape for being the nearest to the given point.
///
Expand All @@ -104,7 +107,7 @@ impl<V: Point, Shape: Bounded<f32, 3>> BvhTraverseDistance<V, Shape> for BvhNode
node_index: usize,
origin: &V,
shapes: &[Shape],
indices: &mut Vec<(usize, f32)>,
best_candidate: &mut (usize, f32),
best_min_distance: &mut f32,
best_max_distance: &mut f32,
) {
Expand All @@ -118,7 +121,7 @@ impl<V: Point, Shape: Bounded<f32, 3>> BvhTraverseDistance<V, Shape> for BvhNode
} => {
// Compute min/max dists for both children
let [child_l_dists, child_r_dists] =
[child_l_aabb, child_r_aabb].map(|aabb| aabb.get_min_max_distances(origin));
[child_l_aabb, child_r_aabb].map(|aabb| aabb.get_min_max_distances_2(origin));

// Update best max distance before traversing children to avoid unnecessary traversals
// where right node prunes left node.
Expand All @@ -129,40 +132,33 @@ impl<V: Point, Shape: Bounded<f32, 3>> BvhTraverseDistance<V, Shape> for BvhNode
(child_l_dists, child_l_index),
(child_r_dists, child_r_index),
] {
// Node is better by a margin.
if dist_max <= *best_min_distance {
indices.clear();
}

// Node might contain a candidate
if dist_min <= *best_max_distance {
Self::nearest_candidates_recursive(
nodes,
*index,
origin,
shapes,
indices,
best_candidate,
best_min_distance,
best_max_distance,
);
}
}
}
Self::Leaf { shape_index, .. } => {
let aabb = shapes[*shape_index].aabb();
let (min_dist, max_dist) = aabb.get_min_max_distances(origin);
let dist = shapes[*shape_index].get_dist_2(origin);

if !indices.is_empty() && max_dist < *best_min_distance {
if dist < *best_min_distance {
// Node is better by a margin
indices.clear();
*best_min_distance = dist;
*best_max_distance = dist;
// we reached a leaf, we add it to the list of indices since it is a potential candidate
*best_candidate = (*shape_index, dist);
} else if dist < *best_max_distance {
// Also update min_dist here since we have a credible (small) bounding box.
*best_max_distance = best_max_distance.min(dist);
}

// Also update min_dist here since we have a credible (small) bounding box.
*best_min_distance = best_min_distance.min(min_dist);
*best_max_distance = best_max_distance.min(max_dist);

// we reached a leaf, we add it to the list of indices since it is a potential candidate
indices.push((*shape_index, min_dist));
}
}
}
Expand Down
103 changes: 63 additions & 40 deletions mesh_to_sdf/src/generate/generic/bvh.rs
View file
Original file line number Diff line number Diff line change
@@ -1,19 +1,21 @@
//! Module containing the `generate_sdf_bvh` function.
use core::cmp::Ordering;

use bvh::{bounding_hierarchy::BoundingHierarchy, bvh::Bvh};
use itertools::Itertools;
use rayon::prelude::*;

use crate::{bvh_ext::BvhDistance, compare_distances, geo, Point, SignMethod, Topology};
use crate::{
bvh_ext::{BvhDistance, BvhNodeExt},
geo, Point, SignMethod, Topology,
};

/// A node in the BVH tree containing the data for a triangle.
/// Public in the crate so that the grid generation can use it.
/// `RtreeBvh` uses its own version of this struct to be able to implement the `RTreeObject` trait.
#[derive(Clone)]
pub struct BvhNode<V: Point> {
pub vertex_indices: (usize, usize, usize),
pub vertices: (V, V, V),
pub node_index: usize,
pub bounding_box: (V, V),
}
Expand Down Expand Up @@ -44,6 +46,15 @@ impl<V: Point> bvh::bounding_hierarchy::BHShape<f32, 3> for BvhNode<V> {
}
}

impl<V: Point> BvhNodeExt<V> for BvhNode<V> {
fn get_dist_2(&self, point: &V) -> f32 {
let a = &self.vertices.0;
let b = &self.vertices.1;
let c = &self.vertices.2;
geo::point_triangle_distance2(point, a, b, c)
}
}

/// Generate a signed distance field from a mesh using a bvh.
/// Query points are expected to be in the same space as the mesh.
///
Expand All @@ -62,6 +73,11 @@ where
let mut bvh_nodes = Topology::get_triangles(vertices, indices)
.map(|triangle| BvhNode {
vertex_indices: triangle,
vertices: (
vertices[triangle.0],
vertices[triangle.1],
vertices[triangle.2],
),
node_index: 0,
bounding_box: geo::triangle_bounding_box(
&vertices[triangle.0],
Expand All @@ -76,33 +92,15 @@ where
query_points
.par_iter()
.map(|point| {
let bvh_indices = bvh.nearest_candidates(point, &bvh_nodes);
let best_candidate = bvh.nearest_candidates(point, &bvh_nodes);

let mut min_dist = f32::MAX;
if sign_method == SignMethod::Normal {
for index in &bvh_indices {
let triangle = &bvh_nodes[*index];
let a = &vertices[triangle.vertex_indices.0];
let b = &vertices[triangle.vertex_indices.1];
let c = &vertices[triangle.vertex_indices.2];
let distance = geo::point_triangle_signed_distance(point, a, b, c);

if compare_distances(min_dist, distance) == Ordering::Greater {
min_dist = distance;
}
}
min_dist
let triangle = &bvh_nodes[best_candidate.0];
let a = &vertices[triangle.vertex_indices.0];
let b = &vertices[triangle.vertex_indices.1];
let c = &vertices[triangle.vertex_indices.2];
geo::point_triangle_signed_distance(point, a, b, c)
} else {
for index in &bvh_indices {
let triangle = &bvh_nodes[*index];
let a = &vertices[triangle.vertex_indices.0];
let b = &vertices[triangle.vertex_indices.1];
let c = &vertices[triangle.vertex_indices.2];
let distance = geo::point_triangle_distance(point, a, b, c);

min_dist = min_dist.min(distance);
}

let alignments = [
(geo::GridAlign::X, nalgebra::Vector3::new(1.0, 0.0, 0.0)),
(geo::GridAlign::Y, nalgebra::Vector3::new(0.0, 1.0, 0.0)),
Expand Down Expand Up @@ -135,9 +133,9 @@ where

// Return inside if at least two are insides.
if insides > 1 {
-min_dist
-best_candidate.1
} else {
min_dist
best_candidate.1
}
}
})
Expand Down Expand Up @@ -233,19 +231,32 @@ mod tests {
SignMethod::Raycast,
);

let mut fails = 0;
// Test against generate_sdf
// TODO: sometimes fails
// thread 'tests::test_generate_bvh_big' panicked at mesh_to_sdf\src\lib.rs:1232:13:
// i: 17435: 0.0076956493 0.030284861
for (i, (sdf, grid_sdf)) in sdf.iter().zip(grid_sdf.iter()).enumerate() {
// Assert we have the same absolute value.
assert!(
(sdf - grid_sdf).abs() < 0.01,
"cell: {:?}: {} {}",
grid.get_cell_integer_coordinates(i),
(sdf.abs() - grid_sdf.abs()).abs() < 0.01,
"i: {}: {} {}",
i,
sdf,
grid_sdf
);

// Count sign issues.
if (sdf - grid_sdf).abs() > 0.01 {
fails += 1;
}
}

// We can expect sign to be different for some cells.
// This is because the rtree only outputs the closest triangle
// while bvh works with a subset of close triangles and refines the sign based on those.
assert!(
(fails as f32 / sdf.len() as f32) < 0.01,
"fails: {fails}/{}",
sdf.len()
);
}

#[test]
Expand Down Expand Up @@ -293,19 +304,31 @@ mod tests {
SignMethod::Normal,
);

let mut fails = 0;
// Test against generate_sdf
for (i, (sdf, grid_sdf)) in sdf.iter().zip(grid_sdf.iter()).enumerate() {
// TODO: sometimes fails
// thread 'tests::test_generate_bvh_big' panicked at mesh_to_sdf\src\lib.rs:1232:13:
// i: 17435: 0.0076956493 0.030284861
// i: 1742: 0.09342232 -0.094851956
// Assert we have the same absolute value.
assert!(
(sdf - grid_sdf).abs() < 0.01,
(sdf.abs() - grid_sdf.abs()).abs() < 0.01,
"i: {}: {} {}",
i,
sdf,
grid_sdf
);

// Count sign issues.
if (sdf - grid_sdf).abs() > 0.01 {
fails += 1;
}
}

// We can expect sign to be different for some cells.
// This is because the rtree only outputs the closest triangle
// while bvh works with a subset of close triangles and refines the sign based on those.
assert!(
(fails as f32 / sdf.len() as f32) < 0.01,
"fails: {fails}/{}",
sdf.len()
);
}
}
Loading

0 comments on commit 18f0bee

Please sign in to comment.