From 24bf8753b206b1cdb179fbcb06a05751a3937165 Mon Sep 17 00:00:00 2001 From: hal8174 Date: Fri, 3 Jan 2025 23:36:18 +0100 Subject: [PATCH] Add Acceleration Structure Interface --- .../src/acceleration_structure/mod.rs | 119 ++++++++++ .../acceleration_structure/triangle_bvh.rs | 222 ++++++++++++++++++ ray-tracing-scene/src/examples/cornell2.rs | 2 +- ray-tracing-scene/src/examples/mod.rs | 5 + .../src/examples/stanford_dragon_as.rs | 126 ++++++++++ ray-tracing-scene/src/lib.rs | 2 + ray-tracing-scene/src/triangle_bvh.rs | 2 +- ray-tracing-scene/src/util.rs | 43 ++++ 8 files changed, 519 insertions(+), 2 deletions(-) create mode 100644 ray-tracing-scene/src/acceleration_structure/mod.rs create mode 100644 ray-tracing-scene/src/acceleration_structure/triangle_bvh.rs create mode 100644 ray-tracing-scene/src/examples/stanford_dragon_as.rs create mode 100644 ray-tracing-scene/src/util.rs diff --git a/ray-tracing-scene/src/acceleration_structure/mod.rs b/ray-tracing-scene/src/acceleration_structure/mod.rs new file mode 100644 index 0000000..a953375 --- /dev/null +++ b/ray-tracing-scene/src/acceleration_structure/mod.rs @@ -0,0 +1,119 @@ +use std::sync::Arc; + +use rand::seq::SliceRandom; +use ray_tracing_core::{ + prelude::*, + scene::{Intersection, LightSample, Scene}, +}; +use sampling::sample_triangle; + +use crate::util::triangle_normal; + +pub mod triangle_bvh; + +pub trait AccelerationStructure { + fn intersect(&self, ray: Ray, min: Float, max: Float) -> Option<(Float, Dir3, T)>; +} + +#[derive(Debug)] +pub struct ASMaterial { + pub material: Option>>, + pub light: Option>>, +} + +impl ASMaterial { + pub fn new_material + 'static>(material: M) -> Self { + Self { + material: Some(Box::new(material)), + light: None, + } + } + pub fn new_light + 'static>(light: L) -> Self { + Self { + material: None, + light: Some(Box::new(light)), + } + } + pub fn new_material_light + 'static, L: Light + 'static>( + material: M, + light: L, + ) -> Self { + Self { + material: Some(Box::new(material)), + light: Some(Box::new(light)), + } + } +} + +#[derive(Debug)] +pub struct AccelerationStructureScene { + acceleration_structure: A, + materials: Arc<[ASMaterial]>, + lights: Vec<(usize, [Pos3; 3])>, +} + +impl Clone for AccelerationStructureScene { + fn clone(&self) -> Self { + Self { + acceleration_structure: self.acceleration_structure.clone(), + materials: Arc::clone(&self.materials), + lights: self.lights.clone(), + } + } +} + +impl AccelerationStructureScene { + pub fn new(a: A, materials: Arc<[ASMaterial]>, lights: Vec<(usize, [Pos3; 3])>) -> Self { + Self { + acceleration_structure: a, + materials, + lights, + } + } +} + +impl, R: Rng> Scene for AccelerationStructureScene { + fn intersect( + &self, + ray: Ray, + min: Float, + max: Float, + ) -> Option> { + let (t, n, i) = self.acceleration_structure.intersect(ray, min, max)?; + + let material = &self.materials[i as usize]; + + Some(Intersection::new( + t, + n, + material.material.as_deref(), + material.light.as_deref(), + )) + } + + fn sample_light( + &self, + _w_in: Dir3, + _intersection: &ray_tracing_core::scene::Intersection<'_, R>, + rng: &mut R, + ) -> Option> { + let t = self.lights.choose(rng); + + if let Some(&(t, v)) = t { + let light = self.materials[t].light.as_ref().unwrap().as_ref(); + let b = sample_triangle(rng.gen()); + + let n = triangle_normal(v); + let area = n.length() * 0.5; + + Some(LightSample::new( + Pos3::from_barycentric(v, b), + 1.0 / ((self.lights.len() as Float) * area), + n.normalize(), + light, + )) + } else { + None + } + } +} diff --git a/ray-tracing-scene/src/acceleration_structure/triangle_bvh.rs b/ray-tracing-scene/src/acceleration_structure/triangle_bvh.rs new file mode 100644 index 0000000..2134fdf --- /dev/null +++ b/ray-tracing-scene/src/acceleration_structure/triangle_bvh.rs @@ -0,0 +1,222 @@ +use ray_tracing_core::prelude::*; + +use crate::util::triangle_normal; + +use super::AccelerationStructure; + +type Index = u32; + +#[derive(Debug, Clone)] +pub struct TriangleBVH { + vertices: Vec, + triangles: Vec>, + bvh: Vec, +} + +#[derive(Debug, Clone, Copy)] +enum Node { + Inner { + left: Index, + left_aabb: AABB, + right: Index, + right_aabb: AABB, + }, + Leaf { + start: Index, + count: Index, + }, +} + +#[derive(Debug, Clone, Copy)] +pub struct Triangle { + pub vertices: [Index; 3], + pub data: T, +} + +impl Triangle { + pub fn new(vertices: [Index; 3], data: T) -> Self { + Self { vertices, data } + } +} + +impl TriangleBVH { + pub fn new(vertices: Vec, mut triangles: Vec>) -> Self { + let mut bvh = vec![Node::Leaf { + start: 0, + count: triangles.len() as Index, + }]; + + let aabb = Self::calculate_aabb(&vertices, &triangles); + + Self::build_bvh(&vertices, &mut triangles, &mut bvh, 0, aabb); + + Self { + vertices, + triangles, + bvh, + } + } + + fn calculate_aabb(vertices: &[Pos3], triangles: &[Triangle]) -> AABB { + let mut aabb = AABB::new( + vertices[triangles[0].vertices[0] as usize], + vertices[triangles[0].vertices[1] as usize], + ); + aabb = aabb.extend(vertices[triangles[0].vertices[2] as usize]); + + for t in triangles.iter().skip(1) { + for v in t.vertices { + aabb = aabb.extend(vertices[v as usize]); + } + } + + aabb + } + + fn build_bvh( + vertices: &[Pos3], + triangles: &mut [Triangle], + bvh: &mut Vec, + node: usize, + aabb: AABB, + ) { + let (start, count) = if let Node::Leaf { start, count } = bvh[node] { + (start, count) + } else { + unreachable!() + }; + + if count < 8 { + return; + } + + let size = aabb.size(); + + let dim = if size.x() > Float::max(size.y(), size.z()) { + 0 + } else if size.y() > Float::max(size.x(), size.z()) { + 1 + } else { + 2 + }; + + let get_key = |t: &Triangle| { + t.vertices + .iter() + .map(|&v| vertices[v as usize][dim]) + .sum::() + / 3.0 + }; + + triangles.sort_by(|a, b| get_key(a).partial_cmp(&get_key(b)).unwrap()); + + let (left_range, right_range) = triangles.split_at_mut((count / 2) as usize); + + let left_aabb = Self::calculate_aabb(vertices, left_range); + let right_aabb = Self::calculate_aabb(vertices, right_range); + + let i = bvh.len() as Index; + bvh[node] = Node::Inner { + left: i, + left_aabb, + right: i + 1, + right_aabb, + }; + + bvh.push(Node::Leaf { + start, + count: left_range.len() as Index, + }); + bvh.push(Node::Leaf { + start: start + left_range.len() as Index, + count: right_range.len() as Index, + }); + + Self::build_bvh(vertices, left_range, bvh, i as usize, left_aabb); + Self::build_bvh(vertices, right_range, bvh, i as usize + 1, right_aabb); + } + + fn intersect_bvh( + &self, + node: Index, + ray: Ray, + min: Float, + max: Float, + ) -> Option<(Float, Dir3, T)> { + match self.bvh[node as usize] { + Node::Inner { + left, + left_aabb, + right, + right_aabb, + } => { + let left_intersect = left_aabb.intersect_ray(ray, min, max); + let right_intersect = right_aabb.intersect_ray(ray, min, max); + + match (left_intersect, right_intersect) { + (None, None) => None, + (None, Some(_)) => self.intersect_bvh(right, ray, min, max), + (Some(_), None) => self.intersect_bvh(left, ray, min, max), + (Some(l), Some(r)) => { + let close; + let far; + if l < r { + close = left; + far = right; + } else { + close = right; + far = left; + } + + if let Some(close_intersect) = self.intersect_bvh(close, ray, min, max) { + if let Some(far_intersect) = self + .intersect_bvh(far, ray, min, Float::min(max, close_intersect.0)) + .filter(|far_intersect| far_intersect.0 < close_intersect.0) + { + Some(far_intersect) + } else { + Some(close_intersect) + } + } else { + self.intersect_bvh(far, ray, min, max) + } + } + } + } + Node::Leaf { start, count } => { + let mut intersection = None; + for i in start..(start + count) { + if let Some(t) = crate::util::triangle_intersection( + ray, + self.get_vertices(&self.triangles[i as usize]), + ) { + if min <= t + && t <= max + && !intersection.is_some_and(|(old_t, _, _)| t > old_t) + { + let n = triangle_normal(self.get_vertices(&self.triangles[i as usize])) + .normalize(); + intersection.replace((t, n, self.triangles[i as usize].data)); + } + } + } + + intersection + } + } + } + + fn get_vertices(&self, triangle: &Triangle) -> [Pos3; 3] { + [ + self.vertices[triangle.vertices[0] as usize], + self.vertices[triangle.vertices[1] as usize], + self.vertices[triangle.vertices[2] as usize], + ] + } +} + +impl AccelerationStructure for TriangleBVH { + fn intersect(&self, ray: Ray, min: Float, max: Float) -> Option<(Float, Dir3, T)> { + self.intersect_bvh(0, ray, min, max) + } +} diff --git a/ray-tracing-scene/src/examples/cornell2.rs b/ray-tracing-scene/src/examples/cornell2.rs index 98b7a37..5536df1 100644 --- a/ray-tracing-scene/src/examples/cornell2.rs +++ b/ray-tracing-scene/src/examples/cornell2.rs @@ -1,7 +1,7 @@ use rand::Rng; use ray_tracing_core::{light::AreaLight, prelude::*, scene::Scene}; use ray_tracing_material::{diffuse::DiffuseMaterial, mirror::Mirror}; -use std::{cell::OnceCell, fmt::Debug}; +use std::cell::OnceCell; use crate::{ parse_obj::ObjData, diff --git a/ray-tracing-scene/src/examples/mod.rs b/ray-tracing-scene/src/examples/mod.rs index d7eae54..96da02e 100644 --- a/ray-tracing-scene/src/examples/mod.rs +++ b/ray-tracing-scene/src/examples/mod.rs @@ -23,9 +23,14 @@ pub fn example_scenes() -> HashMap<&'static str, Box { + scene: OnceCell, R>>, +} + +impl StanfordDragon { + pub fn new() -> Self { + Self { + scene: OnceCell::new(), + } + } +} + +impl ExampleScene for StanfordDragon { + fn get_scene(&self) -> Box + Sync> { + let s = self.scene.get_or_init(|| { + let obj = ObjData::new("ray-tracing-scene/obj/stanford_dragon.obj").unwrap(); + + let mut triangles = obj + .triangle_faces() + .map(|t| Triangle::new(t.v, 0)) + .collect::>(); + + let color = Color::new(0.2, 0.2, 0.9); + let materials = vec![ + ASMaterial::new_material(Microfacet::new(BeckmannDistribution::new(0.01), color)), + ASMaterial::new_material(OrenNayar::new(0.5, Color::new(0.8, 0.8, 0.8))), + ASMaterial::new_material(OrenNayar::new(0.5, Color::new(0.9, 0.0, 0.0))), + ASMaterial::new_material(OrenNayar::new(0.5, Color::new(0.0, 0.9, 0.0))), + ASMaterial::new_light(AreaLight::new(Color::white() * 30.0)), + ]; + + let mut vertices = obj.vertices; + for v in &mut vertices { + *v = Pos3::new(v.x(), v.z() + 60.0, -v.y()); + } + + let side_length = 400.0; + let light_size = 150.0; + let light_offset = 0.01; + let offset = vertices.len() as u32; + vertices.extend_from_slice(&[ + Pos3::new(side_length, 2.0 * side_length, side_length), + Pos3::new(side_length, 2.0 * side_length, -side_length), + Pos3::new(side_length, 0.0, side_length), + Pos3::new(side_length, 0.0, -side_length), + Pos3::new(-side_length, 2.0 * side_length, side_length), + Pos3::new(-side_length, 2.0 * side_length, -side_length), + Pos3::new(-side_length, 0.0, side_length), + Pos3::new(-side_length, 0.0, -side_length), + Pos3::new(light_size, 2.0 * side_length - light_offset, light_size), + Pos3::new(light_size, 2.0 * side_length - light_offset, -light_size), + Pos3::new(-light_size, 2.0 * side_length - light_offset, light_size), + Pos3::new(-light_size, 2.0 * side_length - light_offset, -light_size), + ]); + + triangles.extend_from_slice(&[ + Triangle::new([offset, 1 + offset, 2 + offset], 1), + Triangle::new([1 + offset, 3 + offset, 2 + offset], 1), + Triangle::new([offset, 4 + offset, 1 + offset], 1), + Triangle::new([1 + offset, 4 + offset, 5 + offset], 1), + Triangle::new([2 + offset, 3 + offset, 6 + offset], 1), + Triangle::new([3 + offset, 7 + offset, 6 + offset], 1), + Triangle::new([offset, 2 + offset, 4 + offset], 2), + Triangle::new([6 + offset, 4 + offset, 2 + offset], 2), + Triangle::new([1 + offset, 5 + offset, 3 + offset], 3), + Triangle::new([7 + offset, 3 + offset, 5 + offset], 3), + Triangle::new([8 + offset, 10 + offset, 9 + offset], 4), + Triangle::new([11 + offset, 9 + offset, 10 + offset], 4), + ]); + + let lights: Vec<_> = triangles + .iter() + .filter(|f| materials[f.data as usize].light.is_some()) + .map(|t| { + ( + t.data as usize, + [ + vertices[t.vertices[0] as usize], + vertices[t.vertices[1] as usize], + vertices[t.vertices[2] as usize], + ], + ) + }) + .collect(); + let a = TriangleBVH::new(vertices, triangles); + AccelerationStructureScene::new(a, materials.into(), lights) + }); + + let n: AccelerationStructureScene<_, _> = s.clone(); + Box::new(n) + } + + fn get_camera_pos(&self) -> Pos3 { + Pos3::new(-150.0, 160.0, 250.0) + } + + fn get_camera_look_at(&self) -> Pos3 { + Pos3::new(0.0, 60.0, 0.0) + } + + fn get_camera_up(&self) -> Dir3 { + Dir3::up() + } + + fn get_horizontal_fov(&self) -> Float { + Float::to_radians(90.0) + } +} diff --git a/ray-tracing-scene/src/lib.rs b/ray-tracing-scene/src/lib.rs index 87fb664..f12a912 100644 --- a/ray-tracing-scene/src/lib.rs +++ b/ray-tracing-scene/src/lib.rs @@ -1,4 +1,6 @@ +pub mod acceleration_structure; pub mod basic_scene; pub mod examples; pub mod parse_obj; pub mod triangle_bvh; +pub mod util; diff --git a/ray-tracing-scene/src/triangle_bvh.rs b/ray-tracing-scene/src/triangle_bvh.rs index c50700a..3a77caa 100644 --- a/ray-tracing-scene/src/triangle_bvh.rs +++ b/ray-tracing-scene/src/triangle_bvh.rs @@ -226,12 +226,12 @@ impl TriangleBVH { count: triangles.len() as Index, }]; let aabb = calculate_aabb(&vertices, &triangles); - build_bvh(&vertices, &mut triangles, &mut bvh, 0, aabb); let lights: Vec<_> = triangles .iter() .filter(|f| materials[f.material as usize].light.is_some()) .map(Triangle::clone) .collect(); + build_bvh(&vertices, &mut triangles, &mut bvh, 0, aabb); Self { vertices, bvh, diff --git a/ray-tracing-scene/src/util.rs b/ray-tracing-scene/src/util.rs new file mode 100644 index 0000000..ca95c11 --- /dev/null +++ b/ray-tracing-scene/src/util.rs @@ -0,0 +1,43 @@ +use ray_tracing_core::prelude::*; + +pub fn triangle_intersection(ray: Ray, v: [Pos3; 3]) -> Option { + let e1 = v[1] - v[0]; + let e2 = v[2] - v[0]; + + let ray_cross_e2 = Dir3::cross(ray.dir(), e2); + let det = e1.dot(ray_cross_e2); + + if det > -f32::EPSILON && det < f32::EPSILON { + return None; // This ray is parallel to this triangle. + } + + let inv_det = 1.0 / det; + let s = ray.start() - v[0]; + let u = inv_det * s.dot(ray_cross_e2); + if !(0.0..=1.0).contains(&u) { + return None; + } + + let s_cross_e1 = s.cross(e1); + let v = inv_det * Dir3::dot(ray.dir(), s_cross_e1); + if v < 0.0 || u + v > 1.0 { + return None; + } + // At this stage we can compute t to find out where the intersection point is on the line. + let t = inv_det * e2.dot(s_cross_e1); + + if t > Float::EPSILON { + // ray intersection + Some(t) + } else { + // This means that there is a line intersection but not a ray intersection. + None + } +} + +pub fn triangle_normal(v: [Pos3; 3]) -> Dir3 { + let e1 = v[1] - v[0]; + let e2 = v[2] - v[0]; + + Dir3::cross(e1, e2) +}