aselimov/roctree
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Improved BVH using surface area heuristic*

Browse Source 
- Add custom error type
- Update more mesh based calculations to f32
This commit is contained in:
Alex Selimov 2025-03-08 23:09:15 -05:00
parent 6f404ba2a1
commit ba74db96f0
5 changed files with 235 additions and 76 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

173
src/bvh.rs
View File

@ -1,4 +1,17 @@
use crate::{primitives::Box3, stl::parser::StlSolid};
use core::f64;
use crate::{errors::RoctreeError, primitives::Box3, stl::parser::StlSolid};
/// Number of potential split planes to check
const NBINS: usize = 8;
/// Struct to hold bin information, this includes an axis-aligned bounding box as well as the count
/// of triangles with the
#[derive(Default, Clone)]
struct Bin {
aabb: Box3,
tri_count: usize,
}
/// A node representing a node single in the bounding volume hierarchy
#[derive(Default)]
@ -10,6 +23,16 @@ struct BvhNode {
prim_count: usize,
}
impl BvhNode {
/// This computes the computational cost associated with the specific node.
/// The cost is an aggregate that considers the size of the bounding box (how likely it will
/// intersect a ray) * the number of triangles (the number of ray-triangle checks we have to
/// perform if an intersect occurs)
fn calculate_node_cost(&self) -> f32 {
self.aabb.surface_area() * self.prim_count as f32
}
}
/// Struct representing the total bounding volume hierarchy
pub struct Bvh {
nodes: Vec<BvhNode>,
@ -17,14 +40,15 @@ pub struct Bvh {
}
/// Return type for get_subdivision_position_and_axis
struct SplitPlane {
struct BestSplitPlane {
axis: usize,
pos: f64,
cost: f64,
}
impl Bvh {
/// Create a new Bounding volume hierarchy from an STL solid.
pub fn new(solid: &StlSolid) -> Self {
pub fn new(solid: &StlSolid) -> Result<Bvh, RoctreeError> {
// Initialize the root node
let mut nodes = vec![BvhNode {
prim_idx: 0,
@ -40,16 +64,22 @@ impl Bvh {
tri_idx: Vec::from_iter(0..solid.triangles.len()),
};
bvh.subdivide(0, solid);
bvh
bvh.subdivide(0, solid)?;
Ok(bvh)
}
/// Subdivide a node, subdivision doesn't occur if the node has only 2 primitives
pub fn subdivide(&mut self, node_idx: usize, solid: &StlSolid) {
let SplitPlane {
pub fn subdivide(&mut self, node_idx: usize, solid: &StlSolid) -> Result<(), RoctreeError> {
let BestSplitPlane {
cost,
axis: split_axis,
pos: split_position,
} = self.get_split_plane(node_idx);
} = self.get_best_split(node_idx, solid)?;
// If the split cost is greater than keeping the node intact just return
if cost > self.nodes[node_idx].calculate_node_cost() as f64 {
return Ok(());
}
let prim_count = self.nodes[node_idx].prim_count;
let prim_idx = self.nodes[node_idx].prim_idx;
@ -74,7 +104,7 @@ impl Bvh {
// If all primitives belong to either the left or right half exit because the parent node
// is a leaf
if left_count == 0 || left_count == prim_count {
return;
return Ok(());
}
// Otherwise split the nodes into left and right child
@ -117,24 +147,117 @@ impl Bvh {
self.nodes[node_idx].right_child = right_child;
// Recurse to divide children
self.subdivide(left_child, solid);
self.subdivide(right_child, solid);
self.subdivide(left_child, solid)?;
self.subdivide(right_child, solid)?;
Ok(())
}
/// Calculate the optimal split plane axis and position.
/// This function should use the surface area heuristic to select the optimal split
fn get_split_plane(&self, node_idx: usize) -> SplitPlane {
// Initial approach just splits halfway along the longest axis
// TODO: Improve this to use the surface area heuristic
let axis = (0..3).fold(0, |acc, a| {
if self.nodes[node_idx].aabb.extents[acc] < self.nodes[node_idx].aabb.extents[a] {
a
} else {
acc
}
});
/// Get the best split plane based on the surface area heuristic
/// This procedure works by binning
fn get_best_split(
&self,
node_idx: usize,
solid: &StlSolid,
) -> Result<BestSplitPlane, RoctreeError> {
// Initialize the best split plane information with bad values so we know if no plane gets
// assigned
let mut best_cost = f64::INFINITY;
let mut best_pos = f64::NAN;
let pos = self.nodes[node_idx].aabb.extents[axis] / 2.0;
SplitPlane { axis, pos }
// Pick the longest axis as the splitting axis
let axis = self.nodes[node_idx]
.aabb
.get_extents()
.iter()
.enumerate()
.fold((0_usize, &f32::NEG_INFINITY), |acc, a| {
if acc.1 > a.1 {
acc
} else {
a
}
})
.0;
let node = &self.nodes[node_idx];
// Get minimum and maximum bounds over the centroids off the triangles. This improves
// the binning compared to doing the binning over the vertices
let (min_bd, max_bd) = self.tri_idx[node.prim_idx..node.prim_idx + node.prim_count]
.iter()
.fold((f64::INFINITY, f64::NEG_INFINITY), |bounds, tri_idx| {
(
bounds
.0
.min(solid.triangle_centroids[*tri_idx][axis].into()),
bounds
.1
.max(solid.triangle_centroids[*tri_idx][axis].into()),
)
});
if min_bd == max_bd {
return Err(RoctreeError::SplitPlaneError);
};
let mut bins = vec![Bin::default(); NBINS];
let scale = NBINS as f64 / (max_bd - min_bd);
// Assign the triangles to each of the bins
for tri_idx in self.tri_idx[node.prim_idx..node.prim_idx + node.prim_count].iter() {
let bin_index = usize::min(
NBINS - 1,
((solid.triangle_centroids[*tri_idx][axis] as f64 - min_bd) * scale) as usize,
);
bins[bin_index].tri_count += 1;
bins[bin_index]
.aabb
.grow(solid.vertices[solid.triangles[*tri_idx][0]]);
bins[bin_index]
.aabb
.grow(solid.vertices[solid.triangles[*tri_idx][1]]);
bins[bin_index]
.aabb
.grow(solid.vertices[solid.triangles[*tri_idx][2]]);
}
// Now get the left and right counts for each split plane (based on the bins)
let mut left_area = [0.0; NBINS - 1];
let mut right_area = [0.0; NBINS - 1];
let mut left_count = [0; NBINS - 1];
let mut right_count = [0; NBINS - 1];
let mut left_box = Box3::default();
let mut right_box = Box3::default();
let mut left_sum = 0;
let mut right_sum = 0;
for plane_idx in 0..NBINS - 1 {
// Calculate the left side of the plane
left_sum += bins[plane_idx].tri_count;
left_count[plane_idx] = left_sum;
left_box.grow_by_other_box(&bins[plane_idx].aabb);
left_area[plane_idx] = left_box.surface_area();
// Calculate the right side of the plane
right_sum += bins[NBINS - 1 - plane_idx].tri_count;
right_count[NBINS - 2 - plane_idx] = right_sum;
right_box.grow_by_other_box(&bins[NBINS - 1 - plane_idx].aabb);
right_area[NBINS - 2 - plane_idx] = right_box.surface_area();
}
let scale = 1.0 / scale;
for plane_idx in 0..NBINS - 1 {
let plane_cost: f64 = (left_count[plane_idx] as f32 * left_area[plane_idx]
+ right_count[plane_idx] as f32 * right_area[plane_idx])
.into();
if plane_cost < best_cost {
best_cost = plane_cost;
best_pos = scale * plane_idx as f64;
}
}
Ok(BestSplitPlane {
axis,
pos: best_pos,
cost: best_cost,
})
}
}

19
src/errors.rs Normal file
View File

@ -0,0 +1,19 @@
use std::{error::Error, fmt::Display};
#[derive(Debug)]
pub enum RoctreeError {
SplitPlaneError,
}
impl Error for RoctreeError {}
impl Display for RoctreeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
RoctreeError::SplitPlaneError => write!(
f,
"Failed to get split plane, centroids overlap along longest axis"
),
}
}
}

1
src/lib.rs
View File

@ -1,4 +1,5 @@
pub mod bvh;
pub mod errors;
pub mod octree;
pub mod primitives;
pub mod stl;

106
src/primitives.rs
View File

@ -1,37 +1,36 @@
use std::collections::BTreeSet;
use vecmath::vec3_sub;
use itertools::Itertools;
use vecmath::{vec3_add, vec3_scale};
#[derive(Default)]
#[derive(Default, Clone)]
pub struct Box3 {
pub center: [f64; 3],
pub extents: [f64; 3],
pub min_bd: [f32; 3],
pub max_bd: [f32; 3],
}
impl Box3 {
/// Create a new box from the bounding values
pub fn new_from_bounds(min_bound: [f32; 3], max_bound: [f32; 3]) -> Self {
let mut center = [0.0; 3];
let mut extents = [0.0; 3];
for (i, (min_val, max_val)) in min_bound.iter().zip(max_bound).enumerate() {
center[i] = ((min_val + max_val) / 2.0) as f64;
extents[i] = (max_val - min_val) as f64;
}
Box3 { center, extents }
pub fn new_from_bounds(min_bd: [f32; 3], max_bd: [f32; 3]) -> Self {
Box3 { min_bd, max_bd }
}
/// This function returns the 8 vertices of an axis-aligned bounding box.
pub fn get_vertices(&self) -> Vec<[f64; 3]> {
let half_extent = vec3_scale(self.extents, 0.5);
// Define all eight corners relative to the center.
[-half_extent[0], half_extent[0]]
.into_iter()
.cartesian_product([-half_extent[1], half_extent[1]])
.cartesian_product([-half_extent[2], half_extent[2]])
.map(|((x, y), z)| vec3_add(self.center, [x, y, z]))
.collect()
/// These are ordered
/// 7----6
/// | /| /|
/// | 4----5
/// | | 3--| 2
/// | |/ |/
/// | 0----1
pub fn get_vertices(&self) -> Vec<[f32; 3]> {
vec![
self.min_bd,
[self.max_bd[0], self.min_bd[1], self.min_bd[2]],
[self.max_bd[0], self.max_bd[1], self.min_bd[2]],
[self.min_bd[0], self.max_bd[1], self.min_bd[2]],
[self.max_bd[0], self.min_bd[1], self.max_bd[2]],
[self.max_bd[0], self.max_bd[1], self.max_bd[2]],
self.max_bd,
[self.min_bd[0], self.max_bd[1], self.max_bd[2]],
]
}
/// Expand a box to fit all primitive inside of it.
@ -43,31 +42,50 @@ impl Box3 {
{
// If the current primitive list is empty, zero bounding box and exit
if primitives.is_empty() {
self.center = [0.0; 3];
self.extents = [0.0; 3]
self.min_bd = [0.0; 3];
self.max_bd = [0.0; 3];
return;
}
// Get all of the unique vertex indices so we don't have to check them multiple times
let unique_vertex_idxs: BTreeSet<&usize> = BTreeSet::from_iter(primitives.iter().flatten());
// Now expand the box
let mut min_bd = [f32::INFINITY; 3];
let mut max_bd = [f32::NEG_INFINITY; 3];
for vertex_idx in primitives.iter().flatten() {
self.grow(vertices[*vertex_idx]);
}
}
for idx in unique_vertex_idxs.into_iter() {
for ((vertex_pos, min_pos), max_pos) in vertices[*idx]
.iter()
.zip(min_bd.iter_mut())
.zip(max_bd.iter_mut())
{
*min_pos = vertex_pos.min(*min_pos);
*max_pos = vertex_pos.max(*max_pos);
}
/// Compute the surface area of a box
pub fn surface_area(&self) -> f32 {
let extents = vec3_sub(self.max_bd, self.min_bd);
extents[0] * extents[1] + extents[1] * extents[2] + extents[2] * extents[0]
}
/// Get the extents of the current box
#[inline(always)]
pub fn get_extents(&self) -> [f32; 3] {
vec3_sub(self.max_bd, self.min_bd)
}
/// Grow the axis aligned bounding box to include a specific vertex
#[inline(always)]
pub fn grow(&mut self, vertex: [f32; 3]) {
for (axis, pos) in vertex.iter().enumerate() {
self.min_bd[axis] = self.min_bd[axis].min(*pos);
self.max_bd[axis] = self.max_bd[axis].min(*pos);
}
}
/// Grow the axis aligned bounding box to include a min and max boundary
/// This is useful when trying to combine two bounding boxes
#[inline(always)]
pub fn grow_by_other_box(&mut self, other: &Self) {
for axis in 0..3 {
self.min_bd[axis] = self.min_bd[axis].min(other.min_bd[axis]);
self.max_bd[axis] = self.max_bd[axis].max(other.max_bd[axis]);
}
}
}
pub struct Triangle {
pub v1: [f64; 3],
pub v2: [f64; 3],
pub v3: [f64; 3],
pub v1: [f32; 3],
pub v2: [f32; 3],
pub v3: [f32; 3],
}

12
src/utilities/intersection_checks.rs
View File

@ -1,5 +1,3 @@
use core::f64;
use vecmath::{vec3_cross, vec3_dot, vec3_sub};
use crate::primitives::{Box3, Triangle};
@ -28,8 +26,8 @@ pub fn tri_aabb_intersect(bx: &Box3, tri: &Triangle) -> bool {
let box_projection = project_pnts(&box_vertices, axis);
let tri_projection = project_pnts(&[tri.v1, tri.v2, tri.v3], axis);
let start = f64::max(box_projection.0, tri_projection.0);
let end = f64::min(box_projection.1, tri_projection.1);
let start = f32::max(box_projection.0, tri_projection.0);
let end = f32::min(box_projection.1, tri_projection.1);
if start < end {
return false;
}
@ -38,9 +36,9 @@ pub fn tri_aabb_intersect(bx: &Box3, tri: &Triangle) -> bool {
}
/// Project points unto an Axis and return the max and min
fn project_pnts(pnts: &[[f64; 3]], axis: [f64; 3]) -> (f64, f64) {
let mut min_proj = f64::INFINITY;
let mut max_proj = f64::NEG_INFINITY;
fn project_pnts(pnts: &[[f32; 3]], axis: [f32; 3]) -> (f32, f32) {
let mut min_proj = f32::INFINITY;
let mut max_proj = f32::NEG_INFINITY;
for pnt in pnts {
let projection = vec3_dot(axis, *pnt);
min_proj = min_proj.min(projection);