diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 99994c24..db23fd5c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -42,9 +42,12 @@ jobs: - name: lint run: cargo clippy - - name: test + - name: test (default) run: cargo test + - name: test (web) + run: cargo test --no-default-features --features="web io_ply tooling" + # - name: gaussian render test # run: cargo run --bin test_gaussian diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 00000000..352a6265 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,5 @@ +{ + "rust-analyzer.linkedProjects": [ + "./Cargo.toml" + ] +} \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index 11772fa7..4e4e77af 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "bevy_gaussian_splatting" description = "bevy gaussian splatting render pipeline plugin" -version = "1.0.0" +version = "2.0.0" edition = "2021" authors = ["mosure "] license = "MIT" @@ -36,10 +36,23 @@ default-run = "viewer" default = [ "io_flexbuffers", "io_ply", - "morph_particles", + + # "packed", + "planar", + + "buffer_storage", + # "buffer_texture", + + # "f32", + "f16", + + "query_select", + "query_sparse", + "sort_radix", "sort_rayon", "sort_std", + "tooling", "viewer", ] @@ -50,37 +63,69 @@ io_bincode2 = ["bincode2", "flate2"] io_flexbuffers = ["flexbuffers"] io_ply = ["ply-rs"] +material_noise = ["noise", "dep:noise"] + morph_particles = [] +noise = [] + +f32 = [] +f16 = ["half"] + +packed = [] +planar = [] + +buffer_storage = [] +buffer_texture = [] + +query_raycast = [] +query_select = [] +query_sparse = ["kd-tree", "query_select"] + sort_radix = [] -sort_rayon = [ - "rayon", -] +sort_rayon = ["rayon"] sort_std = [] -tooling = [ - "byte-unit", -] +tooling = ["byte-unit"] viewer = [ "bevy-inspector-egui", "bevy_panorbit_camera", + # "bevy_transform_gizmo", +] + +web = [ + "buffer_texture", + "f16", + "io_flexbuffers", + "planar", + "sort_std", + "viewer", + "webgl2", ] +webgl2 = ["bevy/webgl2"] + [dependencies] -bevy-inspector-egui = { version = "0.21", optional = true } -bevy_panorbit_camera = { version = "0.9", optional = true } +bevy-inspector-egui = { version = "0.22", optional = true } +bevy_mod_picking = { version = "0.17", optional = true } +bevy_panorbit_camera = { version = "0.10", optional = true } +bevy_transform_gizmo = { version = "0.9", optional = true } bincode2 = { version = "2.0", optional = true } byte-unit = { version = "5.0", optional = true } bytemuck = "1.14" flate2 = { version = "1.0", optional = true } flexbuffers = { version = "2.0", optional = true } +half = { version = "2.3.1", optional = true, features = ["serde"] } +kd-tree = { version = "0.5", optional = true } +noise = { version = "0.8.2", optional = true } ply-rs = { version = "0.1", optional = true } rand = "0.8" rayon = { version = "1.8", optional = true } serde = "1.0" static_assertions = "1.1" +typenum = "1.17.0" wgpu = "0.17.1" @@ -89,7 +134,6 @@ console_error_panic_hook = "0.1" wasm-bindgen = "0.2" -# TODO: use minimal bevy features [dependencies.bevy] version = "0.12" default-features = false @@ -102,7 +146,7 @@ features = [ [dependencies.web-sys] -version = "0.3.4" +version = "0.3" features = [ 'Document', 'Element', @@ -162,6 +206,10 @@ name = "test_radix" path = "tests/gpu/radix.rs" required-features = ["debug_gpu", "sort_radix"] +[[example]] +name = "minimal" +path = "examples/minimal.rs" + [[bench]] name = "io" diff --git a/README.md b/README.md index a98116aa..7615ec03 100644 --- a/README.md +++ b/README.md @@ -23,10 +23,11 @@ bevy gaussian splatting render pipeline plugin. view the [live demo](https://mos - [X] gaussian cloud particle effects - [X] wasm support /w [live demo](https://mosure.github.io/bevy_gaussian_splatting/index.html?arg1=cactus.gcloud) - [X] depth colorization +- [X] f16 and f32 gcloud +- [X] wgl2 and webgpu - [ ] 4D gaussian cloud wavelet compression - [ ] accelerated spatial queries - [ ] temporal depth sorting -- [ ] f16 and f32 gcloud support - [ ] skeletons - [ ] volume masks - [ ] level of detail @@ -87,9 +88,11 @@ fn setup_gaussian_cloud( - [dreamgaussian](https://github.com/dreamgaussian/dreamgaussian) - [dynamic-3d-gaussians](https://github.com/JonathonLuiten/Dynamic3DGaussians) - [ewa splatting](https://www.cs.umd.edu/~zwicker/publications/EWASplatting-TVCG02.pdf) +- [gaussian-grouping](https://github.com/lkeab/gaussian-grouping) - [gaussian-splatting](https://github.com/graphdeco-inria/gaussian-splatting) - [gaussian-splatting-viewer](https://github.com/limacv/GaussianSplattingViewer/tree/main) - [gaussian-splatting-web](https://github.com/cvlab-epfl/gaussian-splatting-web) +- [gir](https://3dgir.github.io/) - [making gaussian splats smaller](https://aras-p.info/blog/2023/09/13/Making-Gaussian-Splats-smaller/) - [masked-spacetime-hashing](https://github.com/masked-spacetime-hashing/msth) - [onesweep](https://arxiv.org/ftp/arxiv/papers/2206/2206.01784.pdf) @@ -97,7 +100,9 @@ fn setup_gaussian_cloud( - [phys-gaussian](https://xpandora.github.io/PhysGaussian/) - [point-visualizer](https://github.com/mosure/point-visualizer) - [rusty-automata](https://github.com/mosure/rusty-automata) +- [spacetime-gaussians](https://github.com/oppo-us-research/SpacetimeGaussians) - [splat](https://github.com/antimatter15/splat) - [splatter](https://github.com/Lichtso/splatter) - [sturdy-dollop](https://github.com/mosure/sturdy-dollop) +- [sugar](https://github.com/Anttwo/SuGaR) - [taichi_3d_gaussian_splatting](https://github.com/wanmeihuali/taichi_3d_gaussian_splatting) diff --git a/assets/scenes/icecream.gcloud b/assets/scenes/icecream.gcloud index 2d4815c5..7f5d3b21 100644 Binary files a/assets/scenes/icecream.gcloud and b/assets/scenes/icecream.gcloud differ diff --git a/examples/minimal.rs b/examples/minimal.rs new file mode 100644 index 00000000..fbfb347f --- /dev/null +++ b/examples/minimal.rs @@ -0,0 +1,6 @@ +// TODO: minimal app + + +fn main() { + println!("Hello, world!"); +} diff --git a/src/gaussian.rs b/src/gaussian.rs deleted file mode 100644 index a7bd33a1..00000000 --- a/src/gaussian.rs +++ /dev/null @@ -1,280 +0,0 @@ -use rand::{ - seq::SliceRandom, - prelude::Distribution, - Rng, -}; -use std::marker::Copy; - -use bevy::{ - prelude::*, - reflect::TypeUuid, - render::render_resource::ShaderType, -}; -use bytemuck::{ - Pod, - Zeroable, -}; -use serde::{ - Deserialize, - Serialize, - Serializer, - ser::SerializeTuple, -}; - -use crate::sort::SortMode; - - -const fn num_sh_coefficients(degree: usize) -> usize { - if degree == 0 { - 1 - } else { - 2 * degree + 1 + num_sh_coefficients(degree - 1) - } -} -const SH_DEGREE: usize = 3; -pub const SH_CHANNELS: usize = 3; -pub const MAX_SH_COEFF_COUNT_PER_CHANNEL: usize = num_sh_coefficients(SH_DEGREE); -pub const MAX_SH_COEFF_COUNT: usize = MAX_SH_COEFF_COUNT_PER_CHANNEL * SH_CHANNELS; -#[derive( - Clone, - Copy, - Debug, - PartialEq, - Reflect, - ShaderType, - Pod, - Zeroable, - Serialize, - Deserialize, -)] -#[repr(C)] -pub struct SphericalHarmonicCoefficients { - #[serde(serialize_with = "coefficients_serializer", deserialize_with = "coefficients_deserializer")] - pub coefficients: [f32; MAX_SH_COEFF_COUNT], -} -impl Default for SphericalHarmonicCoefficients { - fn default() -> Self { - Self { - coefficients: [0.0; MAX_SH_COEFF_COUNT], - } - } -} -fn coefficients_serializer(n: &[f32; MAX_SH_COEFF_COUNT], s: S) -> Result -where - S: Serializer, -{ - let mut tup = s.serialize_tuple(MAX_SH_COEFF_COUNT)?; - for &x in n.iter() { - tup.serialize_element(&x)?; - } - - tup.end() -} - -fn coefficients_deserializer<'de, D>(d: D) -> Result<[f32; MAX_SH_COEFF_COUNT], D::Error> -where - D: serde::Deserializer<'de>, -{ - struct CoefficientsVisitor; - - impl<'de> serde::de::Visitor<'de> for CoefficientsVisitor { - type Value = [f32; MAX_SH_COEFF_COUNT]; - - fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { - formatter.write_str("an array of floats") - } - - fn visit_seq(self, mut seq: A) -> Result<[f32; MAX_SH_COEFF_COUNT], A::Error> - where - A: serde::de::SeqAccess<'de>, - { - let mut coefficients = [0.0; MAX_SH_COEFF_COUNT]; - - for (i, coefficient) in coefficients.iter_mut().enumerate().take(MAX_SH_COEFF_COUNT) { - *coefficient = seq - .next_element()? - .ok_or_else(|| serde::de::Error::invalid_length(i, &self))?; - } - Ok(coefficients) - } - } - - d.deserialize_tuple(MAX_SH_COEFF_COUNT, CoefficientsVisitor) -} - - -#[derive( - Clone, - Debug, - Default, - Copy, - PartialEq, - Reflect, - ShaderType, - Pod, - Zeroable, - Serialize, - Deserialize, -)] -#[repr(C)] -// TODO: support f16 gaussian clouds (shader and asset loader) -pub struct Gaussian { - pub rotation: [f32; 4], - pub position: [f32; 4], - pub scale_opacity: [f32; 4], - pub spherical_harmonic: SphericalHarmonicCoefficients, -} - - -#[derive( - Asset, - Clone, - Debug, - Default, - PartialEq, - Reflect, - TypeUuid, - Serialize, - Deserialize, -)] -#[uuid = "ac2f08eb-bc32-aabb-ff21-51571ea332d5"] -pub struct GaussianCloud { - pub gaussians: Vec, -} - -impl GaussianCloud { - pub fn test_model() -> Self { - let origin = Gaussian { - rotation: [ - 1.0, - 0.0, - 0.0, - 0.0, - ], - position: [ - 0.0, - 0.0, - 0.0, - 1.0, - ], - scale_opacity: [ - 0.5, - 0.5, - 0.5, - 0.5, - ], - spherical_harmonic: SphericalHarmonicCoefficients { - coefficients: [ - 1.0, 0.0, 1.0, - 0.0, 0.5, 0.0, - 0.3, 0.2, 0.0, - 0.4, 0.0, 0.2, - 0.1, 0.0, 0.0, - 0.0, 0.3, 0.3, - 0.0, 1.0, 1.0, - 0.3, 0.0, 0.0, - 0.0, 0.0, 0.0, - 0.0, 0.3, 1.0, - 0.5, 0.3, 0.0, - 0.2, 0.3, 0.1, - 0.6, 0.3, 0.1, - 0.0, 0.3, 0.2, - 0.0, 0.5, 0.3, - 0.6, 0.1, 0.2, - ], - }, - }; - let mut cloud = GaussianCloud { - gaussians: Vec::new(), - ..default() - }; - - for &x in [-0.5, 0.5].iter() { - for &y in [-0.5, 0.5].iter() { - for &z in [-0.5, 0.5].iter() { - let mut g = origin; - g.position = [x, y, z, 1.0]; - cloud.gaussians.push(g); - - let mut rng = rand::thread_rng(); - cloud.gaussians.last_mut().unwrap().spherical_harmonic.coefficients.shuffle(&mut rng); - } - } - } - - cloud.gaussians.push(cloud.gaussians[0]); - - cloud - } -} - - -#[derive(Component, Reflect, Clone)] -#[reflect(Component)] -pub struct GaussianCloudSettings { - pub aabb: bool, - pub global_scale: f32, - pub global_transform: GlobalTransform, - pub visualize_bounding_box: bool, - pub visualize_depth: bool, - pub sort_mode: SortMode, -} - -impl Default for GaussianCloudSettings { - fn default() -> Self { - Self { - aabb: false, - global_scale: 2.0, - global_transform: Transform::IDENTITY.into(), - visualize_bounding_box: false, - visualize_depth: false, - sort_mode: SortMode::default(), - } - } -} - -impl Distribution for rand::distributions::Standard { - fn sample(&self, rng: &mut R) -> Gaussian { - Gaussian { - rotation: [ - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - ], - position: [ - rng.gen_range(-20.0..20.0), - rng.gen_range(-20.0..20.0), - rng.gen_range(-20.0..20.0), - rng.gen_range(-1.0..1.0), - ], - scale_opacity: [ - rng.gen_range(0.0..1.0), - rng.gen_range(0.0..1.0), - rng.gen_range(0.0..1.0), - rng.gen_range(0.0..0.8), - ], - spherical_harmonic: SphericalHarmonicCoefficients { - coefficients: { - let mut coefficients = [0.0; MAX_SH_COEFF_COUNT]; - for coefficient in coefficients.iter_mut() { - *coefficient = rng.gen_range(-1.0..1.0); - } - coefficients - }, - }, - } - } -} - -pub fn random_gaussians(n: usize) -> GaussianCloud { - let mut rng = rand::thread_rng(); - let mut gaussians = Vec::with_capacity(n); - for _ in 0..n { - gaussians.push(rng.gen()); - } - GaussianCloud { - gaussians, - ..default() - } -} diff --git a/src/gaussian/cloud.rs b/src/gaussian/cloud.rs new file mode 100644 index 00000000..b8ad4b92 --- /dev/null +++ b/src/gaussian/cloud.rs @@ -0,0 +1,415 @@ +use rand::{ + seq::SliceRandom, + Rng, +}; +use std::iter::FromIterator; + +use bevy::{ + prelude::*, + reflect::TypeUuid, +}; +use serde::{ + Deserialize, + Serialize, +}; + +#[cfg(feature = "sort_rayon")] +use rayon::prelude::*; + +#[allow(unused_imports)] +use crate::{ + gaussian::{ + f32::{ + Position, + PositionVisibility, + Rotation, + ScaleOpacity, + }, + packed::Gaussian, + }, + material::spherical_harmonics::{ + HALF_SH_COEFF_COUNT, + SH_COEFF_COUNT, + SphericalHarmonicCoefficients, + }, +}; + +#[cfg(feature = "f16")] +use crate::gaussian::f16::{ + RotationScaleOpacityPacked128, + pack_f32s_to_u32, + }; + + +#[derive( + Asset, + Clone, + Debug, + Default, + PartialEq, + Reflect, + TypeUuid, + Serialize, + Deserialize, +)] +#[uuid = "ac2f08eb-bc32-aabb-ff21-51571ea332d5"] +pub struct GaussianCloud { + pub position_visibility: Vec, + + pub spherical_harmonic: Vec, + + #[cfg(feature = "f16")] + pub rotation_scale_opacity_packed128: Vec, + + #[cfg(feature = "f32")] + pub rotation: Vec, + #[cfg(feature = "f32")] + pub scale_opacity: Vec, +} + +impl GaussianCloud { + pub fn is_empty(&self) -> bool { + self.position_visibility.is_empty() + } + + pub fn len(&self) -> usize { + self.position_visibility.len() + } + + pub fn len_sqrt_ceil(&self) -> usize { + (self.len() as f32).sqrt().ceil() as usize + } + + pub fn square_len(&self) -> usize { + self.len_sqrt_ceil().pow(2) + } + + pub fn position(&self, index: usize) -> &[f32; 3] { + &self.position_visibility[index].position + } + + pub fn position_mut(&mut self, index: usize) -> &mut [f32; 3] { + &mut self.position_visibility[index].position + } + + pub fn position_iter(&self) -> impl Iterator + '_ { + self.position_visibility.iter() + .map(|position_visibility| &position_visibility.position) + } + + #[cfg(feature = "sort_rayon")] + pub fn position_par_iter(&self) -> impl IndexedParallelIterator { + self.position_visibility.par_iter() + .map(|position_visibility| &position_visibility.position) + } + + + pub fn visibility(&self, index: usize) -> f32 { + self.position_visibility[index].visibility + } + + pub fn visibility_mut(&mut self, index: usize) -> &mut f32 { + &mut self.position_visibility[index].visibility + } + + + // pub fn rotation(&self, index: usize) -> &[f32; 4] { + // #[cfg(feature = "f16")] + // return &self.rotation_scale_opacity_packed128[index].rotation; + + // #[cfg(feature = "f32")] + // return &self.rotation[index].rotation; + // } + + // pub fn rotation_mut(&mut self, index: usize) -> &mut [f32; 4] { + // #[cfg(feature = "f16")] + // return &mut self.rotation_scale_opacity_packed128[index].rotation; + + // #[cfg(feature = "f32")] + // return &mut self.rotation[index].rotation; + // } + + + // pub fn scale(&self, index: usize) -> &[f32; 3] { + // #[cfg(feature = "f16")] + // return &self.rotation_scale_opacity_packed128[index].scale; + + // #[cfg(feature = "f32")] + // return &self.scale_opacity[index].scale; + // } + + // pub fn scale_mut(&mut self, index: usize) -> &mut [f32; 3] { + // #[cfg(feature = "f16")] + // return &mut self.rotation_scale_opacity_packed128[index].scale; + + // #[cfg(feature = "f32")] + // return &mut self.scale_opacity[index].scale; + // } + + #[cfg(feature = "f16")] + pub fn gaussian(&self, index: usize) -> Gaussian { + let rso = self.rotation_scale_opacity_packed128[index]; + + let rotation = rso.rotation(); + let scale_opacity = rso.scale_opacity(); + + Gaussian { + position_visibility: self.position_visibility[index], + spherical_harmonic: self.spherical_harmonic[index], + rotation, + scale_opacity, + } + } + + #[cfg(feature = "f32")] + pub fn gaussian(&self, index: usize) -> Gaussian { + Gaussian { + position_visibility: self.position_visibility[index], + spherical_harmonic: self.spherical_harmonic[index], + rotation: self.rotation[index], + scale_opacity: self.scale_opacity[index], + } + } + + #[cfg(feature = "f16")] + pub fn gaussian_iter(&self) -> impl Iterator + '_ { + self.position_visibility.iter() + .zip(self.spherical_harmonic.iter()) + .zip(self.rotation_scale_opacity_packed128.iter()) + .map(|((position_visibility, spherical_harmonic), rotation_scale_opacity)| { + Gaussian { + position_visibility: *position_visibility, + spherical_harmonic: *spherical_harmonic, + + rotation: rotation_scale_opacity.rotation(), + scale_opacity: rotation_scale_opacity.scale_opacity(), + } + }) + } + + #[cfg(feature = "f32")] + pub fn gaussian_iter(&self) -> impl Iterator + '_ { + self.position_visibility.iter() + .zip(self.spherical_harmonic.iter()) + .zip(self.rotation.iter()) + .zip(self.scale_opacity.iter()) + .map(|(((position_visibility, spherical_harmonic), rotation), scale_opacity)| { + Gaussian { + position_visibility: *position_visibility, + spherical_harmonic: *spherical_harmonic, + + rotation: *rotation, + scale_opacity: *scale_opacity, + } + }) + } + + + pub fn spherical_harmonic(&self, index: usize) -> &SphericalHarmonicCoefficients { + &self.spherical_harmonic[index] + } + + pub fn spherical_harmonic_mut(&mut self, index: usize) -> &mut SphericalHarmonicCoefficients { + &mut self.spherical_harmonic[index] + } + + pub fn resize_to_square(&mut self) { + #[cfg(all(feature = "buffer_texture", feature = "f16"))] + { + self.position_visibility.resize(self.square_len(), PositionVisibility::default()); + self.spherical_harmonic.resize(self.square_len(), SphericalHarmonicCoefficients::default()); + self.rotation_scale_opacity_packed128.resize(self.square_len(), RotationScaleOpacityPacked128::default()); + } + + #[cfg(all(feature = "buffer_texture", feature = "f32"))] + { + self.position_visibility.resize(self.square_len(), PositionVisibility::default()); + self.spherical_harmonic.resize(self.square_len(), SphericalHarmonicCoefficients::default()); + self.rotation.resize(self.square_len(), Rotation::default()); + self.scale_opacity.resize(self.square_len(), ScaleOpacity::default()); + } + } +} + + +impl GaussianCloud { + #[cfg(feature = "f16")] + pub fn subset(&self, indicies: &[usize]) -> Self { + let mut position_visibility = Vec::with_capacity(indicies.len()); + let mut spherical_harmonic = Vec::with_capacity(indicies.len()); + let mut rotation_scale_opacity_packed128 = Vec::with_capacity(indicies.len()); + + for &index in indicies.iter() { + position_visibility.push(self.position_visibility[index]); + spherical_harmonic.push(self.spherical_harmonic[index]); + rotation_scale_opacity_packed128.push(self.rotation_scale_opacity_packed128[index]); + } + + Self { + position_visibility, + spherical_harmonic, + rotation_scale_opacity_packed128, + } + } + + #[cfg(feature = "f32")] + pub fn subset(&self, indicies: &[usize]) -> Self { + let mut position_visibility = Vec::with_capacity(indicies.len()); + let mut spherical_harmonic = Vec::with_capacity(indicies.len()); + let mut rotation = Vec::with_capacity(indicies.len()); + let mut scale_opacity = Vec::with_capacity(indicies.len()); + + for &index in indicies.iter() { + position_visibility.push(self.position_visibility[index]); + spherical_harmonic.push(self.spherical_harmonic[index]); + rotation.push(self.rotation[index]); + scale_opacity.push(self.scale_opacity[index]); + } + + Self { + position_visibility, + spherical_harmonic, + rotation, + scale_opacity, + } + } + + #[cfg(feature = "f32")] + pub fn to_packed(&self) -> Vec { + let mut gaussians = Vec::with_capacity(self.len()); + + for index in 0..self.len() { + gaussians.push(self.gaussian(index)); + } + + gaussians + } +} + + +impl GaussianCloud { + #[cfg(feature = "f16")] + pub fn from_gaussians(gaussians: Vec) -> Self { + let mut position_visibility = Vec::with_capacity(gaussians.len()); + let mut spherical_harmonic = Vec::with_capacity(gaussians.len()); + let mut rotation_scale_opacity_packed128 = Vec::with_capacity(gaussians.len()); + + for gaussian in gaussians { + position_visibility.push(gaussian.position_visibility); + spherical_harmonic.push(gaussian.spherical_harmonic); + + rotation_scale_opacity_packed128.push(RotationScaleOpacityPacked128::from_gaussian(&gaussian)); + } + + #[allow(unused_mut)] + let mut cloud = GaussianCloud { + position_visibility, + spherical_harmonic, + rotation_scale_opacity_packed128, + }; + + cloud.resize_to_square(); + + cloud + } + + #[cfg(feature = "f32")] + pub fn from_gaussians(gaussians: Vec) -> Self { + let mut position_visibility = Vec::with_capacity(gaussians.len()); + let mut spherical_harmonic = Vec::with_capacity(gaussians.len()); + let mut rotation = Vec::with_capacity(gaussians.len()); + let mut scale_opacity = Vec::with_capacity(gaussians.len()); + + for gaussian in gaussians { + position_visibility.push(gaussian.position_visibility); + spherical_harmonic.push(gaussian.spherical_harmonic); + + rotation.push(gaussian.rotation); + scale_opacity.push(gaussian.scale_opacity); + } + + Self { + position_visibility, + spherical_harmonic, + rotation, + scale_opacity, + } + } + + pub fn test_model() -> Self { + let mut rng = rand::thread_rng(); + + let origin = Gaussian { + rotation: [ + 1.0, + 0.0, + 0.0, + 0.0, + ].into(), + position_visibility: [ + 0.0, + 0.0, + 0.0, + 1.0, + ].into(), + scale_opacity: [ + 0.5, + 0.5, + 0.5, + 0.5, + ].into(), + spherical_harmonic: SphericalHarmonicCoefficients { + coefficients: { + #[cfg(feature = "f16")] + { + let mut coefficients = [0 as u32; HALF_SH_COEFF_COUNT]; + + for coefficient in coefficients.iter_mut() { + let upper = rng.gen_range(-1.0..1.0); + let lower = rng.gen_range(-1.0..1.0); + + *coefficient = pack_f32s_to_u32(upper, lower); + } + + coefficients + } + + #[cfg(feature = "f32")] + { + let mut coefficients = [0.0; SH_COEFF_COUNT]; + + for coefficient in coefficients.iter_mut() { + *coefficient = rng.gen_range(-1.0..1.0); + } + + coefficients + } + }, + }, + }; + let mut gaussians: Vec = Vec::new(); + + for &x in [-0.5, 0.5].iter() { + for &y in [-0.5, 0.5].iter() { + for &z in [-0.5, 0.5].iter() { + let mut g = origin; + g.position_visibility = [x, y, z, 1.0].into(); + gaussians.push(g); + + gaussians.last_mut().unwrap().spherical_harmonic.coefficients.shuffle(&mut rng); + } + } + } + + gaussians.push(gaussians[0]); + + GaussianCloud::from_gaussians(gaussians) + } +} + +impl FromIterator for GaussianCloud { + fn from_iter>(iter: I) -> Self { + let gaussians = iter.into_iter().collect::>(); + GaussianCloud::from_gaussians(gaussians) + } +} diff --git a/src/gaussian/f16.rs b/src/gaussian/f16.rs new file mode 100644 index 00000000..ced8d2a5 --- /dev/null +++ b/src/gaussian/f16.rs @@ -0,0 +1,153 @@ +use std::marker::Copy; + +use half::f16; + +use bevy::{ + prelude::*, + render::render_resource::ShaderType, +}; +use bytemuck::{ + Pod, + Zeroable, +}; +use serde::{ + Deserialize, + Serialize, +}; + +use crate::gaussian::{ + f32::{ + Rotation, + ScaleOpacity, + }, + packed::Gaussian, +}; + + +#[derive( + Clone, + Debug, + Default, + Copy, + PartialEq, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct RotationScaleOpacityPacked128 { + #[reflect(ignore)] + pub rotation: [u32; 2], + #[reflect(ignore)] + pub scale_opacity: [u32; 2], +} + +impl RotationScaleOpacityPacked128 { + pub fn from_gaussian(gaussian: &Gaussian) -> Self { + Self { + rotation: [ + pack_f32s_to_u32(gaussian.rotation.rotation[0], gaussian.rotation.rotation[1]), + pack_f32s_to_u32(gaussian.rotation.rotation[2], gaussian.rotation.rotation[3]), + ], + scale_opacity: [ + pack_f32s_to_u32(gaussian.scale_opacity.scale[0], gaussian.scale_opacity.scale[1]), + pack_f32s_to_u32(gaussian.scale_opacity.scale[2], gaussian.scale_opacity.opacity), + ], + } + } + + pub fn rotation(&self) -> Rotation { + let (u0, l0) = unpack_u32_to_f32s(self.rotation[0]); + let (u1, l1) = unpack_u32_to_f32s(self.rotation[1]); + + Rotation { + rotation: [ + u0, + l0, + u1, + l1, + ], + } + } + + pub fn scale_opacity(&self) -> ScaleOpacity { + let (u0, l0) = unpack_u32_to_f32s(self.scale_opacity[0]); + let (u1, l1) = unpack_u32_to_f32s(self.scale_opacity[1]); + + ScaleOpacity { + scale: [ + u0, + l0, + u1, + ], + opacity: l1, + } + } +} + +impl From<[f32; 8]> for RotationScaleOpacityPacked128 { + fn from(rotation_scale_opacity: [f32; 8]) -> Self { + Self { + rotation: [ + pack_f32s_to_u32(rotation_scale_opacity[0], rotation_scale_opacity[1]), + pack_f32s_to_u32(rotation_scale_opacity[2], rotation_scale_opacity[3]), + ], + scale_opacity: [ + pack_f32s_to_u32(rotation_scale_opacity[4], rotation_scale_opacity[5]), + pack_f32s_to_u32(rotation_scale_opacity[6], rotation_scale_opacity[7]), + ], + } + } +} + +impl From<[f16; 8]> for RotationScaleOpacityPacked128 { + fn from(rotation_scale_opacity: [f16; 8]) -> Self { + Self { + rotation: [ + pack_f16s_to_u32(rotation_scale_opacity[0], rotation_scale_opacity[1]), + pack_f16s_to_u32(rotation_scale_opacity[2], rotation_scale_opacity[3]), + ], + scale_opacity: [ + pack_f16s_to_u32(rotation_scale_opacity[4], rotation_scale_opacity[5]), + pack_f16s_to_u32(rotation_scale_opacity[6], rotation_scale_opacity[7]), + ], + } + } +} + +impl From<[u32; 4]> for RotationScaleOpacityPacked128 { + fn from(rotation_scale_opacity: [u32; 4]) -> Self { + Self { + rotation: [rotation_scale_opacity[0], rotation_scale_opacity[1]], + scale_opacity: [rotation_scale_opacity[2], rotation_scale_opacity[3]], + } + } +} + + +pub fn pack_f32s_to_u32(upper: f32, lower: f32) -> u32 { + pack_f16s_to_u32( + f16::from_f32(upper), + f16::from_f32(lower), + ) +} + +pub fn pack_f16s_to_u32(upper: f16, lower: f16) -> u32 { + let upper_bits = (upper.to_bits() as u32) << 16; + let lower_bits = lower.to_bits() as u32; + upper_bits | lower_bits +} + +pub fn unpack_u32_to_f16s(value: u32) -> (f16, f16) { + let upper = f16::from_bits((value >> 16) as u16); + let lower = f16::from_bits((value & 0xFFFF) as u16); + (upper, lower) +} + +pub fn unpack_u32_to_f32s(value: u32) -> (f32, f32) { + let (upper, lower) = unpack_u32_to_f16s(value); + (upper.to_f32(), lower.to_f32()) +} diff --git a/src/gaussian/f32.rs b/src/gaussian/f32.rs new file mode 100644 index 00000000..f9a1c1da --- /dev/null +++ b/src/gaussian/f32.rs @@ -0,0 +1,105 @@ +use std::marker::Copy; + +use bevy::{ + prelude::*, + render::render_resource::ShaderType, +}; +use bytemuck::{ + Pod, + Zeroable, +}; +use serde::{ + Deserialize, + Serialize, +}; + + +pub type Position = [f32; 3]; + +#[derive( + Clone, + Debug, + Default, + Copy, + PartialEq, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct PositionVisibility { + pub position: Position, + pub visibility: f32, +} + +impl From<[f32; 4]> for PositionVisibility { + fn from(position_visibility: [f32; 4]) -> Self { + Self { + position: [ + position_visibility[0], + position_visibility[1], + position_visibility[2], + ], + visibility: position_visibility[3], + } + } +} + +#[derive( + Clone, + Debug, + Default, + Copy, + PartialEq, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct Rotation { + pub rotation: [f32; 4], +} + +impl From<[f32; 4]> for Rotation { + fn from(rotation: [f32; 4]) -> Self { + Self { rotation } + } +} + +#[derive( + Clone, + Debug, + Default, + Copy, + PartialEq, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct ScaleOpacity { + pub scale: [f32; 3], + pub opacity: f32, +} + +impl From<[f32; 4]> for ScaleOpacity { + fn from(scale_opacity: [f32; 4]) -> Self { + Self { + scale: [ + scale_opacity[0], + scale_opacity[1], + scale_opacity[2], + ], + opacity: scale_opacity[3], + } + } +} diff --git a/src/gaussian/group.rs b/src/gaussian/group.rs new file mode 100644 index 00000000..e69de29b diff --git a/src/gaussian/mod.rs b/src/gaussian/mod.rs new file mode 100644 index 00000000..4053137b --- /dev/null +++ b/src/gaussian/mod.rs @@ -0,0 +1,37 @@ +use static_assertions::assert_cfg; + +pub mod cloud; +pub mod f32; +pub mod packed; +pub mod rand; +pub mod settings; + +#[cfg(feature = "f16")] +pub mod f16; + + +assert_cfg!( + any( + feature = "f16", + feature = "f32", + ), + "specify one of the following features: f16, f32", +); + +assert_cfg!( + any( + feature = "packed", + feature = "planar", + ), + "specify one of the following features: packed, planar", +); + + +// PACKED_f16 is not supported +assert_cfg!( + not(all( + feature = "f16", + feature = "packed", + )), + "f16 and packed are incompatible", +); diff --git a/src/gaussian/packed.rs b/src/gaussian/packed.rs new file mode 100644 index 00000000..3c9af38e --- /dev/null +++ b/src/gaussian/packed.rs @@ -0,0 +1,41 @@ +use std::marker::Copy; + +use bevy::prelude::*; +use bytemuck::{ + Pod, + Zeroable, +}; +use serde::{ + Deserialize, + Serialize, +}; + +use crate::{ + gaussian::f32::{ + PositionVisibility, + Rotation, + ScaleOpacity, + }, + material::spherical_harmonics::SphericalHarmonicCoefficients, +}; + + +#[derive( + Clone, + Debug, + Default, + Copy, + PartialEq, + Reflect, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct Gaussian { + pub rotation: Rotation, + pub position_visibility: PositionVisibility, + pub scale_opacity: ScaleOpacity, + pub spherical_harmonic: SphericalHarmonicCoefficients, +} diff --git a/src/gaussian/rand.rs b/src/gaussian/rand.rs new file mode 100644 index 00000000..4d4cab4c --- /dev/null +++ b/src/gaussian/rand.rs @@ -0,0 +1,82 @@ +use rand::{ + prelude::Distribution, + Rng, +}; + +#[cfg(feature = "f16")] +use crate::gaussian::f16::pack_f32s_to_u32; + +#[allow(unused_imports)] +use crate::{ + gaussian::{ + cloud::GaussianCloud, + packed::Gaussian, + }, + material::spherical_harmonics::{ + HALF_SH_COEFF_COUNT, + SH_COEFF_COUNT, + SphericalHarmonicCoefficients, + }, +}; + + +impl Distribution for rand::distributions::Standard { + fn sample(&self, rng: &mut R) -> Gaussian { + Gaussian { + rotation: [ + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + ].into(), + position_visibility: [ + rng.gen_range(-20.0..20.0), + rng.gen_range(-20.0..20.0), + rng.gen_range(-20.0..20.0), + 1.0, + ].into(), + scale_opacity: [ + rng.gen_range(0.0..1.0), + rng.gen_range(0.0..1.0), + rng.gen_range(0.0..1.0), + rng.gen_range(0.0..0.8), + ].into(), + spherical_harmonic: SphericalHarmonicCoefficients { + coefficients: { + #[cfg(feature = "f16")] + { + let mut coefficients: [u32; HALF_SH_COEFF_COUNT] = [0; HALF_SH_COEFF_COUNT]; + for coefficient in coefficients.iter_mut() { + let upper = rng.gen_range(-1.0..1.0); + let lower = rng.gen_range(-1.0..1.0); + + *coefficient = pack_f32s_to_u32(upper, lower); + } + coefficients + } + + #[cfg(feature = "f32")] + { + let mut coefficients = [0.0; SH_COEFF_COUNT]; + for coefficient in coefficients.iter_mut() { + *coefficient = rng.gen_range(-1.0..1.0); + } + coefficients + } + }, + }, + } + } +} + +pub fn random_gaussians(n: usize) -> GaussianCloud { + let mut rng = rand::thread_rng(); + let mut gaussians: Vec = Vec::with_capacity(n); + + for _ in 0..n { + gaussians.push(rng.gen()); + } + + GaussianCloud::from_gaussians(gaussians) +} + diff --git a/src/gaussian/settings.rs b/src/gaussian/settings.rs new file mode 100644 index 00000000..6f4d7b66 --- /dev/null +++ b/src/gaussian/settings.rs @@ -0,0 +1,47 @@ +use bevy::prelude::*; + +use crate::sort::SortMode; + + +#[derive( + Clone, + Copy, + Debug, + Default, + Eq, + Hash, + PartialEq, + Reflect, +)] +pub enum GaussianCloudDrawMode { + #[default] + All, + Selected, + HighlightSelected, +} + +#[derive(Component, Reflect, Clone)] +#[reflect(Component)] +pub struct GaussianCloudSettings { + pub aabb: bool, + pub global_scale: f32, + pub global_transform: GlobalTransform, + pub visualize_bounding_box: bool, + pub visualize_depth: bool, + pub sort_mode: SortMode, + pub draw_mode: GaussianCloudDrawMode, +} + +impl Default for GaussianCloudSettings { + fn default() -> Self { + Self { + aabb: false, + global_scale: 2.0, + global_transform: Transform::IDENTITY.into(), + visualize_bounding_box: false, + visualize_depth: false, + sort_mode: SortMode::default(), + draw_mode: GaussianCloudDrawMode::default(), + } + } +} diff --git a/src/io/loader.rs b/src/io/loader.rs index 9bf22d0f..21e6c086 100644 --- a/src/io/loader.rs +++ b/src/io/loader.rs @@ -1,3 +1,4 @@ +#[allow(unused_imports)] use std::io::{ BufReader, Cursor, @@ -47,12 +48,8 @@ impl AssetLoader for GaussianCloudLoader { let mut f = BufReader::new(cursor); let gaussians = crate::io::ply::parse_ply(&mut f)?; - let cloud = GaussianCloud { - gaussians, - ..Default::default() - }; - Ok(cloud) + Ok(GaussianCloud::from_gaussians(gaussians)) } #[cfg(not(feature = "io_ply"))] diff --git a/src/io/mod.rs b/src/io/mod.rs index cfe1eb24..5eb3c54f 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -1,6 +1,7 @@ pub mod codec; pub mod gcloud; pub mod loader; +pub mod writer; #[cfg(feature = "io_ply")] pub mod ply; diff --git a/src/io/ply.rs b/src/io/ply.rs index 5db6ef1b..e9d2a277 100644 --- a/src/io/ply.rs +++ b/src/io/ply.rs @@ -8,10 +8,13 @@ use ply_rs::{ parser::Parser, }; -use crate::gaussian::{ - Gaussian, - MAX_SH_COEFF_COUNT_PER_CHANNEL, - SH_CHANNELS, +use crate::{ + material::spherical_harmonics::{ + SH_CHANNELS, + SH_COEFF_COUNT, + SH_COEFF_COUNT_PER_CHANNEL, + }, + gaussian::packed::Gaussian, }; @@ -24,28 +27,42 @@ impl PropertyAccess for Gaussian { fn set_property(&mut self, key: String, property: Property) { match (key.as_ref(), property) { - ("x", Property::Float(v)) => self.position[0] = v, - ("y", Property::Float(v)) => self.position[1] = v, - ("z", Property::Float(v)) => self.position[2] = v, - ("f_dc_0", Property::Float(v)) => self.spherical_harmonic.coefficients[0] = v, - ("f_dc_1", Property::Float(v)) => self.spherical_harmonic.coefficients[1] = v, - ("f_dc_2", Property::Float(v)) => self.spherical_harmonic.coefficients[2] = v, - ("scale_0", Property::Float(v)) => self.scale_opacity[0] = v, - ("scale_1", Property::Float(v)) => self.scale_opacity[1] = v, - ("scale_2", Property::Float(v)) => self.scale_opacity[2] = v, - ("opacity", Property::Float(v)) => self.scale_opacity[3] = 1.0 / (1.0 + (-v).exp()), - ("rot_0", Property::Float(v)) => self.rotation[0] = v, - ("rot_1", Property::Float(v)) => self.rotation[1] = v, - ("rot_2", Property::Float(v)) => self.rotation[2] = v, - ("rot_3", Property::Float(v)) => self.rotation[3] = v, + ("x", Property::Float(v)) => self.position_visibility.position[0] = v, + ("y", Property::Float(v)) => self.position_visibility.position[1] = v, + ("z", Property::Float(v)) => self.position_visibility.position[2] = v, + ("f_dc_0", Property::Float(v)) => self.spherical_harmonic.set(0, v), + ("f_dc_1", Property::Float(v)) => self.spherical_harmonic.set(1, v), + ("f_dc_2", Property::Float(v)) => self.spherical_harmonic.set(2, v), + ("scale_0", Property::Float(v)) => self.scale_opacity.scale[0] = v, + ("scale_1", Property::Float(v)) => self.scale_opacity.scale[1] = v, + ("scale_2", Property::Float(v)) => self.scale_opacity.scale[2] = v, + ("opacity", Property::Float(v)) => self.scale_opacity.opacity = 1.0 / (1.0 + (-v).exp()), + ("rot_0", Property::Float(v)) => self.rotation.rotation[0] = v, + ("rot_1", Property::Float(v)) => self.rotation.rotation[1] = v, + ("rot_2", Property::Float(v)) => self.rotation.rotation[2] = v, + ("rot_3", Property::Float(v)) => self.rotation.rotation[3] = v, (_, Property::Float(v)) if key.starts_with("f_rest_") => { let i = key[7..].parse::().unwrap(); - match i { - _ if i + 3 < self.spherical_harmonic.coefficients.len() => { - self.spherical_harmonic.coefficients[i + 3] = v; - }, - _ => { }, + // interleaved + // if (i + 3) < SH_COEFF_COUNT { + // self.spherical_harmonic.coefficients[i + 3] = v; + // } + + // planar + let channel = i / SH_COEFF_COUNT_PER_CHANNEL; + let coefficient = if SH_COEFF_COUNT_PER_CHANNEL == 1 { + 1 + } else { + (i % (SH_COEFF_COUNT_PER_CHANNEL - 1)) + 1 + }; + + let interleaved_idx = coefficient * SH_CHANNELS + channel; + + if interleaved_idx < SH_COEFF_COUNT { + self.spherical_harmonic.set(interleaved_idx, v); + } else { + // TODO: convert higher degree SH to lower degree SH } } (_, _) => {}, @@ -66,30 +83,15 @@ pub fn parse_ply(mut reader: &mut dyn BufRead) -> Result, std::io: } for gaussian in &mut cloud { - gaussian.position[3] = 1.0; + gaussian.position_visibility.visibility = 1.0; - let mean_scale = (gaussian.scale_opacity[0] + gaussian.scale_opacity[1] + gaussian.scale_opacity[2]) / 3.0; + let mean_scale = (gaussian.scale_opacity.scale[0] + gaussian.scale_opacity.scale[1] + gaussian.scale_opacity.scale[2]) / 3.0; for i in 0..3 { - gaussian.scale_opacity[i] = gaussian.scale_opacity[i] + gaussian.scale_opacity.scale[i] = gaussian.scale_opacity.scale[i] .max(mean_scale - MAX_SIZE_VARIANCE) .min(mean_scale + MAX_SIZE_VARIANCE) .exp(); } - - let sh_src = gaussian.spherical_harmonic.coefficients; - let sh = &mut gaussian.spherical_harmonic.coefficients; - - for (i, sh_src) in sh_src.iter().enumerate().skip(SH_CHANNELS) { - let j = i - SH_CHANNELS; - - let channel = j / (MAX_SH_COEFF_COUNT_PER_CHANNEL - 1); - let coefficient = (j % (MAX_SH_COEFF_COUNT_PER_CHANNEL - 1)) + 1; - - let interleaved_idx = coefficient * SH_CHANNELS + channel; - assert!(interleaved_idx >= SH_CHANNELS); - - sh[interleaved_idx] = *sh_src; - } } Ok(cloud) diff --git a/src/io/writer.rs b/src/io/writer.rs new file mode 100644 index 00000000..e12bd533 --- /dev/null +++ b/src/io/writer.rs @@ -0,0 +1,18 @@ +use std::io::Write; + +use crate::{ + GaussianCloud, + io::codec::GaussianCloudCodec, +}; + + +pub fn write_gaussian_cloud_to_file( + cloud: &GaussianCloud, + path: &str, +) { + let gcloud_file = std::fs::File::create(path).expect("failed to create file"); + let mut gcloud_writer = std::io::BufWriter::new(gcloud_file); + + let data = cloud.encode(); + gcloud_writer.write_all(data.as_slice()).expect("failed to write to gcloud file"); +} diff --git a/src/lib.rs b/src/lib.rs index 6ff57acd..7d363d1d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,24 +1,30 @@ use bevy::prelude::*; pub use gaussian::{ - Gaussian, - GaussianCloud, - GaussianCloudSettings, - SphericalHarmonicCoefficients, - random_gaussians, + packed::Gaussian, + cloud::GaussianCloud, + rand::random_gaussians, + settings::GaussianCloudSettings, }; +pub use material::spherical_harmonics::SphericalHarmonicCoefficients; + use io::loader::GaussianCloudLoader; use render::RenderPipelinePlugin; pub mod gaussian; pub mod io; +pub mod material; pub mod morph; +pub mod query; pub mod render; pub mod sort; pub mod utils; +#[cfg(feature = "noise")] +pub mod noise; + #[derive(Bundle, Default, Reflect)] pub struct GaussianSplattingBundle { @@ -49,6 +55,11 @@ impl Plugin for GaussianSplattingPlugin { app.add_plugins(( RenderPipelinePlugin, + material::MaterialPlugin, + query::QueryPlugin, )); + + #[cfg(feature = "noise")] + app.add_plugins(noise::NoisePlugin); } } diff --git a/src/lighting/environmental.rs b/src/lighting/environmental.rs new file mode 100644 index 00000000..e69de29b diff --git a/src/lighting/mod.rs b/src/lighting/mod.rs new file mode 100644 index 00000000..e69de29b diff --git a/src/material/depth.rs b/src/material/depth.rs new file mode 100644 index 00000000..87e911f1 --- /dev/null +++ b/src/material/depth.rs @@ -0,0 +1,21 @@ +use bevy::{ + prelude::*, + asset::load_internal_asset, +}; + + +const DEPTH_SHADER_HANDLE: Handle = Handle::weak_from_u128(51234253); + + +pub struct DepthMaterialPlugin; + +impl Plugin for DepthMaterialPlugin { + fn build(&self, app: &mut App) { + load_internal_asset!( + app, + DEPTH_SHADER_HANDLE, + "depth.wgsl", + Shader::from_wgsl + ); + } +} diff --git a/src/render/color.wgsl b/src/material/depth.wgsl similarity index 87% rename from src/render/color.wgsl rename to src/material/depth.wgsl index d640dc55..fab6b6a0 100644 --- a/src/render/color.wgsl +++ b/src/material/depth.wgsl @@ -1,4 +1,4 @@ -#define_import_path bevy_gaussian_splatting::color +#define_import_path bevy_gaussian_splatting::depth fn depth_to_rgb(depth: f32, min_depth: f32, max_depth: f32) -> vec3 { diff --git a/src/material/mod.rs b/src/material/mod.rs new file mode 100644 index 00000000..368720f5 --- /dev/null +++ b/src/material/mod.rs @@ -0,0 +1,24 @@ +use bevy::prelude::*; + +pub mod depth; +pub mod spherical_harmonics; + +#[cfg(feature = "material_noise")] +pub mod noise; + + +#[derive(Default)] +pub struct MaterialPlugin; + +impl Plugin for MaterialPlugin { + #[allow(unused)] + fn build(&self, app: &mut App) { + #[cfg(feature = "material_noise")] + app.add_plugins(noise::NoiseMaterialPlugin); + + app.add_plugins(( + depth::DepthMaterialPlugin, + spherical_harmonics::SphericalHarmonicCoefficientsPlugin, + )); + } +} diff --git a/src/material/noise.rs b/src/material/noise.rs new file mode 100644 index 00000000..720a8b94 --- /dev/null +++ b/src/material/noise.rs @@ -0,0 +1,80 @@ +use bevy::prelude::*; +use noise::{ + NoiseFn, + RidgedMulti, + Simplex, +}; + +use crate::{ + Gaussian, + GaussianCloud, +}; + + +#[derive(Component, Debug, Reflect)] +pub struct NoiseMaterial { + pub scale: f32, +} +impl Default for NoiseMaterial { + fn default() -> Self { + NoiseMaterial { + scale: 1.0, + } + } +} + + +#[derive(Default)] +pub struct NoiseMaterialPlugin; + +impl Plugin for NoiseMaterialPlugin { + #[allow(unused)] + fn build(&self, app: &mut App) { + app.register_type::(); + + app.add_systems(Update, apply_noise_cpu); + } +} + + +fn apply_noise_cpu( + mut gaussian_clouds_res: ResMut>, + mut selections: Query<( + Entity, + &Handle, + &NoiseMaterial, + Changed, + )>, +) { + for ( + _entity, + cloud_handle, + noise_material, + changed, + ) in selections.iter_mut() { + if !changed { + continue; + } + + let mut rigid_multi = RidgedMulti::::default(); + rigid_multi.frequency = noise_material.scale as f64; + + let cloud = gaussian_clouds_res.get_mut(cloud_handle).unwrap(); + + cloud.gaussians.iter_mut() + .for_each(|gaussian| { + let point = |gaussian: &Gaussian, idx| { + let x = gaussian.position_visibility[0]; + let y = gaussian.position_visibility[1]; + let z = gaussian.position_visibility[2]; + + [x as f64, y as f64, z as f64, idx as f64] + }; + + for i in 0..gaussian.spherical_harmonic.coefficients.len() { + let noise = rigid_multi.get(point(&gaussian, i)); + gaussian.spherical_harmonic.coefficients[i] = noise as f32; + } + }); + } +} diff --git a/src/material/noise.wgsl b/src/material/noise.wgsl new file mode 100644 index 00000000..e69de29b diff --git a/src/material/pbr.rs b/src/material/pbr.rs new file mode 100644 index 00000000..e69de29b diff --git a/src/material/pbr.wgsl b/src/material/pbr.wgsl new file mode 100644 index 00000000..e69de29b diff --git a/src/material/spherical_harmonics.rs b/src/material/spherical_harmonics.rs new file mode 100644 index 00000000..2db78d43 --- /dev/null +++ b/src/material/spherical_harmonics.rs @@ -0,0 +1,236 @@ +use std::marker::Copy; + +use bevy::{ + prelude::*, + asset::load_internal_asset, + render::render_resource::ShaderType, +}; +use bytemuck::{ + Pod, + Zeroable, +}; +use serde::{ + Deserialize, + Serialize, + Serializer, + ser::SerializeTuple, +}; + +#[cfg(feature = "f16")] +use half::f16; + + +const SPHERICAL_HARMONICS_SHADER_HANDLE: Handle = Handle::weak_from_u128(834667312); + + +pub struct SphericalHarmonicCoefficientsPlugin; + +impl Plugin for SphericalHarmonicCoefficientsPlugin { + fn build(&self, app: &mut App) { + load_internal_asset!( + app, + SPHERICAL_HARMONICS_SHADER_HANDLE, + "spherical_harmonics.wgsl", + Shader::from_wgsl + ); + } +} + + +const fn num_sh_coefficients(degree: usize) -> usize { + if degree == 0 { + 1 + } else { + 2 * degree + 1 + num_sh_coefficients(degree - 1) + } +} + + +#[cfg(feature = "web")] +const SH_DEGREE: usize = 0; + +#[cfg(not(feature = "web"))] +const SH_DEGREE: usize = 3; + +pub const SH_CHANNELS: usize = 3; +pub const SH_COEFF_COUNT_PER_CHANNEL: usize = num_sh_coefficients(SH_DEGREE); +pub const SH_COEFF_COUNT: usize = (SH_COEFF_COUNT_PER_CHANNEL * SH_CHANNELS + 3) & !3; + +pub const HALF_SH_COEFF_COUNT: usize = SH_COEFF_COUNT / 2; +pub const PADDED_HALF_SH_COEFF_COUNT: usize = (HALF_SH_COEFF_COUNT + 3) & !3; + +#[cfg(feature = "f16")] +pub const SH_VEC4_PLANES: usize = PADDED_HALF_SH_COEFF_COUNT / 4; +#[cfg(feature = "f32")] +pub const SH_VEC4_PLANES: usize = SH_COEFF_COUNT / 4; + + +#[cfg(feature = "f16")] +#[derive( + Clone, + Copy, + Debug, + PartialEq, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct SphericalHarmonicCoefficients { + #[reflect(ignore)] + #[serde(serialize_with = "coefficients_serializer", deserialize_with = "coefficients_deserializer")] + pub coefficients: [u32; HALF_SH_COEFF_COUNT], +} + + +#[cfg(feature = "f32")] +#[derive( + Clone, + Copy, + Debug, + PartialEq, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct SphericalHarmonicCoefficients { + #[serde(serialize_with = "coefficients_serializer", deserialize_with = "coefficients_deserializer")] + pub coefficients: [f32; SH_COEFF_COUNT], +} + + +#[cfg(feature = "f16")] +impl Default for SphericalHarmonicCoefficients { + fn default() -> Self { + Self { + coefficients: [0; HALF_SH_COEFF_COUNT], + } + } +} + +#[cfg(feature = "f32")] +impl Default for SphericalHarmonicCoefficients { + fn default() -> Self { + Self { + coefficients: [0.0; SH_COEFF_COUNT], + } + } +} + + +impl SphericalHarmonicCoefficients { + #[cfg(feature = "f16")] + pub fn set(&mut self, index: usize, value: f32) { + let quantized = f16::from_f32(value).to_bits(); + self.coefficients[index / 2] = match index % 2 { + 0 => (self.coefficients[index / 2] & 0xffff0000) | (quantized as u32), + 1 => (self.coefficients[index / 2] & 0x0000ffff) | ((quantized as u32) << 16), + _ => unreachable!(), + }; + } + + #[cfg(feature = "f32")] + pub fn set(&mut self, index: usize, value: f32) { + self.coefficients[index] = value; + } +} + + + +#[cfg(feature = "f16")] +fn coefficients_serializer(n: &[u32; HALF_SH_COEFF_COUNT], s: S) -> Result +where + S: Serializer, +{ + let mut tup = s.serialize_tuple(SH_COEFF_COUNT)?; + for &x in n.iter() { + tup.serialize_element(&x)?; + } + + tup.end() +} + +#[cfg(feature = "f16")] +fn coefficients_deserializer<'de, D>(d: D) -> Result<[u32; HALF_SH_COEFF_COUNT], D::Error> +where + D: serde::Deserializer<'de>, +{ + struct CoefficientsVisitor; + + impl<'de> serde::de::Visitor<'de> for CoefficientsVisitor { + type Value = [u32; HALF_SH_COEFF_COUNT]; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("an array of floats") + } + + fn visit_seq(self, mut seq: A) -> Result<[u32; HALF_SH_COEFF_COUNT], A::Error> + where + A: serde::de::SeqAccess<'de>, + { + let mut coefficients = [0; HALF_SH_COEFF_COUNT]; + + for (i, coefficient) in coefficients.iter_mut().enumerate().take(SH_COEFF_COUNT) { + *coefficient = seq + .next_element()? + .ok_or_else(|| serde::de::Error::invalid_length(i, &self))?; + } + Ok(coefficients) + } + } + + d.deserialize_tuple(SH_COEFF_COUNT, CoefficientsVisitor) +} + + +#[cfg(feature = "f32")] +fn coefficients_serializer(n: &[f32; SH_COEFF_COUNT], s: S) -> Result +where + S: Serializer, +{ + let mut tup = s.serialize_tuple(SH_COEFF_COUNT)?; + for &x in n.iter() { + tup.serialize_element(&x)?; + } + + tup.end() +} + +#[cfg(feature = "f32")] +fn coefficients_deserializer<'de, D>(d: D) -> Result<[f32; SH_COEFF_COUNT], D::Error> +where + D: serde::Deserializer<'de>, +{ + struct CoefficientsVisitor; + + impl<'de> serde::de::Visitor<'de> for CoefficientsVisitor { + type Value = [f32; SH_COEFF_COUNT]; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("an array of floats") + } + + fn visit_seq(self, mut seq: A) -> Result<[f32; SH_COEFF_COUNT], A::Error> + where + A: serde::de::SeqAccess<'de>, + { + let mut coefficients = [0.0; SH_COEFF_COUNT]; + + for (i, coefficient) in coefficients.iter_mut().enumerate().take(SH_COEFF_COUNT) { + *coefficient = seq + .next_element()? + .ok_or_else(|| serde::de::Error::invalid_length(i, &self))?; + } + Ok(coefficients) + } + } + + d.deserialize_tuple(SH_COEFF_COUNT, CoefficientsVisitor) +} diff --git a/src/render/spherical_harmonics.wgsl b/src/material/spherical_harmonics.wgsl similarity index 93% rename from src/render/spherical_harmonics.wgsl rename to src/material/spherical_harmonics.wgsl index 2226d5ca..fbe5465d 100644 --- a/src/render/spherical_harmonics.wgsl +++ b/src/material/spherical_harmonics.wgsl @@ -34,23 +34,28 @@ fn srgb_to_linear(srgb_color: vec3) -> vec3 { fn spherical_harmonics_lookup( ray_direction: vec3, - sh: array, + sh: array, ) -> vec3 { - var rds = ray_direction * ray_direction; + let rds = ray_direction * ray_direction; var color = vec3(0.5); color += shc[ 0] * vec3(sh[0], sh[1], sh[2]); +#if SH_COEFF_COUNT > 11 color += shc[ 1] * vec3(sh[ 3], sh[ 4], sh[ 5]) * ray_direction.y; color += shc[ 2] * vec3(sh[ 6], sh[ 7], sh[ 8]) * ray_direction.z; color += shc[ 3] * vec3(sh[ 9], sh[10], sh[11]) * ray_direction.x; +#endif +#if SH_COEFF_COUNT > 26 color += shc[ 4] * vec3(sh[12], sh[13], sh[14]) * ray_direction.x * ray_direction.y; color += shc[ 5] * vec3(sh[15], sh[16], sh[17]) * ray_direction.y * ray_direction.z; color += shc[ 6] * vec3(sh[18], sh[19], sh[20]) * (2.0 * rds.z - rds.x - rds.y); color += shc[ 7] * vec3(sh[21], sh[22], sh[23]) * ray_direction.x * ray_direction.z; color += shc[ 8] * vec3(sh[24], sh[25], sh[26]) * (rds.x - rds.y); +#endif +#if SH_COEFF_COUNT > 47 color += shc[ 9] * vec3(sh[27], sh[28], sh[29]) * ray_direction.y * (3.0 * rds.x - rds.y); color += shc[10] * vec3(sh[30], sh[31], sh[32]) * ray_direction.x * ray_direction.y * ray_direction.z; color += shc[11] * vec3(sh[33], sh[34], sh[35]) * ray_direction.y * (4.0 * rds.z - rds.x - rds.y); @@ -58,6 +63,7 @@ fn spherical_harmonics_lookup( color += shc[13] * vec3(sh[39], sh[40], sh[41]) * ray_direction.x * (4.0 * rds.z - rds.x - rds.y); color += shc[14] * vec3(sh[42], sh[43], sh[44]) * ray_direction.z * (rds.x - rds.y); color += shc[15] * vec3(sh[45], sh[46], sh[47]) * ray_direction.x * (rds.x - 3.0 * rds.y); +#endif return srgb_to_linear(color); } diff --git a/src/morph/particle.rs b/src/morph/particle.rs index 9a68d875..988790fb 100644 --- a/src/morph/particle.rs +++ b/src/morph/particle.rs @@ -24,7 +24,15 @@ use bevy::{ RenderAssets, RenderAssetPlugin, }, - render_resource::*, + render_resource::{ + Buffer, + BufferInitDescriptor, + BufferUsages, + Extent3d, + ShaderType, + TextureDimension, + TextureFormat, + }, renderer::{ RenderContext, RenderDevice, @@ -53,6 +61,7 @@ use serde::{ use crate::render::{ GaussianCloudBindGroup, GaussianCloudPipeline, + GaussianCloudPipelineKey, GaussianUniformBindGroups, GaussianViewBindGroup, shader_defs, @@ -201,7 +210,7 @@ impl FromWorld for ParticleBehaviorPipeline { ], }); - let shader_defs = shader_defs(false, false, false); + let shader_defs = shader_defs(GaussianCloudPipelineKey::default()); let pipeline_cache = render_world.resource::(); let particle_behavior_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { diff --git a/src/morph/particle.wgsl b/src/morph/particle.wgsl index 002d48bb..fce7bff7 100644 --- a/src/morph/particle.wgsl +++ b/src/morph/particle.wgsl @@ -37,7 +37,7 @@ fn apply_particle_behaviors( let delta_velocity = behavior.acceleration * globals.delta_time + 0.5 * behavior.jerk * globals.delta_time * globals.delta_time; let delta_acceleration = behavior.jerk * globals.delta_time; - let new_position = point.position + delta_position; + let new_position = point.position_visibility + delta_position; let new_velocity = behavior.velocity + delta_velocity; let new_acceleration = behavior.acceleration + delta_acceleration; @@ -47,7 +47,7 @@ fn apply_particle_behaviors( return; } - points[point_index].position = new_position; + points[point_index].position_visibility = new_position; particle_behaviors[behavior_index].velocity = new_velocity; particle_behaviors[behavior_index].acceleration = new_acceleration; } diff --git a/src/noise/mod.rs b/src/noise/mod.rs new file mode 100644 index 00000000..a773e872 --- /dev/null +++ b/src/noise/mod.rs @@ -0,0 +1,22 @@ +use bevy::{ + asset::load_internal_asset, + prelude::*, +}; + + +const NOISE_SHADER_HANDLE: Handle = Handle::weak_from_u128(125722721); + + +#[derive(Default)] +pub struct NoisePlugin; + +impl Plugin for NoisePlugin { + fn build(&self, app: &mut App) { + load_internal_asset!( + app, + NOISE_SHADER_HANDLE, + "noise.wgsl", + Shader::from_wgsl + ); + } +} diff --git a/src/noise/noise.wgsl b/src/noise/noise.wgsl new file mode 100644 index 00000000..70524ba1 --- /dev/null +++ b/src/noise/noise.wgsl @@ -0,0 +1,202 @@ +#define_import_path bevy_gaussian_splatting::noise + + +// MIT License. © Ian McEwan, Stefan Gustavson, Munrocket, Johan Helsing +// +fn mod289_2d(x: vec2) -> vec2 { + return x - floor(x * (1.0 / 289.0)) * 289.0; +} + +fn mod289_3d(x: vec3) -> vec3 { + return x - floor(x * (1.0 / 289.0)) * 289.0; +} + +fn permute_3d(x: vec3) -> vec3 { + return mod289_3d(((x * 34.0) + 1.0) * x); +} + +// MIT License. © Ian McEwan, Stefan Gustavson, Munrocket +fn simplex_2d(v: vec2) -> f32 { + let C = vec4( + 0.211324865405187, // (3.0-sqrt(3.0))/6.0 + 0.366025403784439, // 0.5*(sqrt(3.0)-1.0) + -0.577350269189626, // -1.0 + 2.0 * C.x + 0.024390243902439 // 1.0 / 41.0 + ); + + // First corner + var i = floor(v + dot(v, C.yy)); + let x0 = v - i + dot(i, C.xx); + + // Other corners + var i1 = select(vec2(0., 1.), vec2(1., 0.), x0.x > x0.y); + + // x0 = x0 - 0.0 + 0.0 * C.xx ; + // x1 = x0 - i1 + 1.0 * C.xx ; + // x2 = x0 - 1.0 + 2.0 * C.xx ; + var x12 = x0.xyxy + C.xxzz; + x12.x = x12.x - i1.x; + x12.y = x12.y - i1.y; + + // Permutations + i = mod289_2d(i); // Avoid truncation effects in permutation + + var p = permute_3d(permute_3d(i.y + vec3(0., i1.y, 1.)) + i.x + vec3(0., i1.x, 1.)); + var m = max(0.5 - vec3(dot(x0, x0), dot(x12.xy, x12.xy), dot(x12.zw, x12.zw)), vec3(0.)); + m *= m; + m *= m; + + // Gradients: 41 points uniformly over a line, mapped onto a diamond. + // The ring size 17*17 = 289 is close to a multiple of 41 (41*7 = 287) + let x = 2. * fract(p * C.www) - 1.; + let h = abs(x) - 0.5; + let ox = floor(x + 0.5); + let a0 = x - ox; + + // Normalize gradients implicitly by scaling m + // Approximation of: m *= inversesqrt( a0*a0 + h*h ); + m *= 1.79284291400159 - 0.85373472095314 * (a0 * a0 + h * h); + + // Compute final noise value at P + let g = vec3(a0.x * x0.x + h.x * x0.y, a0.yz * x12.xz + h.yz * x12.yw); + return 130. * dot(m, g); +} + + +// MIT License. © Stefan Gustavson, Munrocket +fn permute4_(x: vec4) -> vec4 { return ((x * 34.0 + 1.0) * x) % vec4(289.0); } +fn taylorInvSqrt4_(r: vec4) -> vec4 { return 1.79284291400159 - 0.85373472095314 * r; } +fn fade3_(t: vec3) -> vec3 { return t * t * t * (t * (t * 6.0 - 15.0) + 10.0); } + +fn perlin_3d(P: vec3) -> f32 { + var Pi0 : vec3 = floor(P); // Integer part for indexing + var Pi1 : vec3 = Pi0 + vec3(1.0); // Integer part + 1 + Pi0 = Pi0 % vec3(289.0); + Pi1 = Pi1 % vec3(289.0); + let Pf0 = fract(P); // Fractional part for interpolation + let Pf1 = Pf0 - vec3(1.0); // Fractional part - 1. + let ix = vec4(Pi0.x, Pi1.x, Pi0.x, Pi1.x); + let iy = vec4(Pi0.yy, Pi1.yy); + let iz0 = Pi0.zzzz; + let iz1 = Pi1.zzzz; + + let ixy = permute4_(permute4_(ix) + iy); + let ixy0 = permute4_(ixy + iz0); + let ixy1 = permute4_(ixy + iz1); + + var gx0: vec4 = ixy0 / 7.0; + var gy0: vec4 = fract(floor(gx0) / 7.0) - 0.5; + gx0 = fract(gx0); + var gz0: vec4 = vec4(0.5) - abs(gx0) - abs(gy0); + var sz0: vec4 = step(gz0, vec4(0.0)); + gx0 = gx0 + sz0 * (step(vec4(0.0), gx0) - 0.5); + gy0 = gy0 + sz0 * (step(vec4(0.0), gy0) - 0.5); + + var gx1: vec4 = ixy1 / 7.0; + var gy1: vec4 = fract(floor(gx1) / 7.0) - 0.5; + gx1 = fract(gx1); + var gz1: vec4 = vec4(0.5) - abs(gx1) - abs(gy1); + var sz1: vec4 = step(gz1, vec4(0.0)); + gx1 = gx1 - sz1 * (step(vec4(0.0), gx1) - 0.5); + gy1 = gy1 - sz1 * (step(vec4(0.0), gy1) - 0.5); + + var g000: vec3 = vec3(gx0.x, gy0.x, gz0.x); + var g100: vec3 = vec3(gx0.y, gy0.y, gz0.y); + var g010: vec3 = vec3(gx0.z, gy0.z, gz0.z); + var g110: vec3 = vec3(gx0.w, gy0.w, gz0.w); + var g001: vec3 = vec3(gx1.x, gy1.x, gz1.x); + var g101: vec3 = vec3(gx1.y, gy1.y, gz1.y); + var g011: vec3 = vec3(gx1.z, gy1.z, gz1.z); + var g111: vec3 = vec3(gx1.w, gy1.w, gz1.w); + + let norm0 = taylorInvSqrt4_( + vec4(dot(g000, g000), dot(g010, g010), dot(g100, g100), dot(g110, g110))); + g000 = g000 * norm0.x; + g010 = g010 * norm0.y; + g100 = g100 * norm0.z; + g110 = g110 * norm0.w; + let norm1 = taylorInvSqrt4_( + vec4(dot(g001, g001), dot(g011, g011), dot(g101, g101), dot(g111, g111))); + g001 = g001 * norm1.x; + g011 = g011 * norm1.y; + g101 = g101 * norm1.z; + g111 = g111 * norm1.w; + + let n000 = dot(g000, Pf0); + let n100 = dot(g100, vec3(Pf1.x, Pf0.yz)); + let n010 = dot(g010, vec3(Pf0.x, Pf1.y, Pf0.z)); + let n110 = dot(g110, vec3(Pf1.xy, Pf0.z)); + let n001 = dot(g001, vec3(Pf0.xy, Pf1.z)); + let n101 = dot(g101, vec3(Pf1.x, Pf0.y, Pf1.z)); + let n011 = dot(g011, vec3(Pf0.x, Pf1.yz)); + let n111 = dot(g111, Pf1); + + var fade_xyz: vec3 = fade3_(Pf0); + let temp = vec4(f32(fade_xyz.z)); // simplify after chrome bug fix + let n_z = mix(vec4(n000, n100, n010, n110), vec4(n001, n101, n011, n111), temp); + let n_yz = mix(n_z.xy, n_z.zw, vec2(f32(fade_xyz.y))); // simplify after chrome bug fix + let n_xyz = mix(n_yz.x, n_yz.y, fade_xyz.x); + return 2.2 * n_xyz; +} + + +// +// by Inigo Quilez +fn hash_23_(p: vec2) -> vec3 { + let q = vec3(dot(p, vec2(127.1, 311.7)), + dot(p, vec2(269.5, 183.3)), + dot(p, vec2(419.2, 371.9))); + return fract(sin(q) * 43758.5453); +} + +fn voro_2d(x: vec2, u: f32, v: f32) -> f32 { + let p = floor(x); + let f = fract(x); + let k = 1.0 + 63.0 * pow(1. - v, 4.0); + var va: f32 = 0.0; + var wt: f32 = 0.0; + for(var j: i32 = -2; j <= 2; j = j + 1) { + for(var i: i32 = -2; i <= 2; i = i + 1) { + let g = vec2(f32(i), f32(j)); + let o = hash_23_(p + g) * vec3(u, u, 1.0); + let r = g - f + o.xy; + let d = dot(r, r); + let ww = pow(1. - smoothstep(0.0, 1.414, sqrt(d)), k); + va = va + o.z * ww; + wt = wt + ww; + } + } + return va / wt; +} + + + +fn nrand(n: vec2) -> f32 { + return fract(sin(dot(n, vec2(12.9898, 4.1414))) * 43758.5453); +} + +fn noise_2d(n: vec2) -> f32 { + let d = vec2(0.0, 1.0); + let b = floor(n); + let f = smoothstep(vec2(0.0), vec2(1.0), fract(n)); + return mix(mix(nrand(b), nrand(b + d.yx), f.x), mix(nrand(b + d.xy), nrand(b + d.yy), f.x), f.y); +} + + +// https://www.shadertoy.com/view/MlVSzw +const ALPHA: f32 = 0.14; +const INV_ALPHA: f32 = 7.14285714286; +const K: f32 = 0.08912676813; + +fn inv_error_function(x: f32) -> f32 { + let y: f32 = log(1.0 - x*x); + let z: f32 = K + 0.5 * y; + return sqrt(sqrt(z*z - y * INV_ALPHA) - z) * sign(x); +} + +// expects n to be in ~[0, 1] +fn gaussian_rand(n: vec2) -> f32 { + let x: f32 = nrand(n * 13.7); + + return inv_error_function(x * 2.0 - 1.0) * 0.3; +} diff --git a/src/query/mod.rs b/src/query/mod.rs new file mode 100644 index 00000000..a885f4ad --- /dev/null +++ b/src/query/mod.rs @@ -0,0 +1,28 @@ +use bevy::prelude::*; + +#[cfg(feature = "query_raycast")] +pub mod raycast; + +#[cfg(feature = "query_select")] +pub mod select; + +#[cfg(feature = "query_sparse")] +pub mod sparse; + + +#[derive(Default)] +pub struct QueryPlugin; + +impl Plugin for QueryPlugin { + #[allow(unused)] + fn build(&self, app: &mut App) { + #[cfg(feature = "query_raycast")] + app.add_plugins(raycast::RaycastSelectionPlugin); + + #[cfg(feature = "query_select")] + app.add_plugins(select::SelectPlugin); + + #[cfg(feature = "query_sparse")] + app.add_plugins(sparse::SparsePlugin); + } +} diff --git a/src/query/raycast.rs b/src/query/raycast.rs new file mode 100644 index 00000000..d3df696c --- /dev/null +++ b/src/query/raycast.rs @@ -0,0 +1,103 @@ +use bevy::{ + prelude::*, + render::mesh::PrimitiveTopology, +}; +use std::collections::BTreeMap; + + +struct Triangle { + vertices: [Vec3; 3], +} + +fn point_in_mesh_system( + mesh_query: Query<(&Handle, &Transform)>, + point_query: Query<(&Point, Entity)>, + meshes: Res>, + mut commands: Commands, +) { + for (mesh_handle, transform) in mesh_query.iter() { + if let Some(mesh) = meshes.get(mesh_handle) { + for (point, entity) in point_query.iter() { + let local_point = transform.compute_matrix().inverse().transform_point3(point.position); + + if is_point_in_mesh(&local_point, mesh) { + commands.entity(entity).insert(InsideMesh); + } + } + } + } +} + +fn is_point_in_mesh(point: &Vec3, mesh: &Mesh) -> bool { + if mesh.primitive_topology() != PrimitiveTopology::TriangleList { + panic!("Mesh must be a triangle list"); + } + + let vertices = if let Some(vertex_attribute) = mesh.attribute(Mesh::ATTRIBUTE_POSITION) { + vertex_attribute + .as_float3() + .expect("Expected vertex positions as Vec3") + } else { + panic!("Mesh does not contain vertex positions"); + }; + + let indices = if let Some(Indices::U32(indices)) = &mesh.indices() { + indices + } else { + panic!("Mesh indices must be of type U32"); + }; + + let mut intersections = 0; + for chunk in indices.chunks_exact(3) { + let triangle = Triangle { + vertices: [ + vertices[chunk[0] as usize], + vertices[chunk[1] as usize], + vertices[chunk[2] as usize], + ], + }; + + let ray_direction = Vec3::new(1.0, 0.0, 0.0); + if ray_intersects_triangle(point, &ray_direction, &triangle) { + intersections += 1; + } + } + + intersections % 2 != 0 +} + + +fn ray_intersects_triangle(ray_origin: &Vec3, ray_direction: &Vec3, triangle: &Triangle) -> bool { + let epsilon = 0.000_001; + let vertex0 = triangle.vertices[0]; + let vertex1 = triangle.vertices[1]; + let vertex2 = triangle.vertices[2]; + + let edge1 = vertex1 - vertex0; + let edge2 = vertex2 - vertex0; + let h = ray_direction.cross(edge2); + let a = edge1.dot(h); + + if a > -epsilon && a < epsilon { + return false; + } + + let f = 1.0 / a; + let s = *ray_origin - vertex0; + let u = f * s.dot(h); + + if u < 0.0 || u > 1.0 { + return false; + } + + let q = s.cross(edge1); + let v = f * ray_direction.dot(q); + + if v < 0.0 || u + v > 1.0 { + return false; + } + + let t = f * edge2.dot(q); + + t > epsilon +} diff --git a/src/query/select.rs b/src/query/select.rs new file mode 100644 index 00000000..015c88dd --- /dev/null +++ b/src/query/select.rs @@ -0,0 +1,178 @@ +use std::iter::FromIterator; + +use bevy::{ + prelude::*, + asset::LoadState, +}; + +use crate::{GaussianCloud, io::writer::write_gaussian_cloud_to_file}; + + +#[derive(Component, Debug, Default, Reflect)] +pub struct Select { + pub indicies: Vec, + pub completed: bool, +} + +impl FromIterator for Select { + fn from_iter>(iter: I) -> Self { + let indicies = iter.into_iter().collect::>(); + Select { indicies, ..Default::default() } + } +} + +impl Select { + pub fn invert(&mut self, cloud_size: usize) -> Select { + let inverted = (0..cloud_size) + .filter(|index| !self.indicies.contains(index)) + .collect::>(); + + Select { + indicies: inverted, + completed: self.completed, + } + } +} + + +#[derive(Default)] +pub struct SelectPlugin; + +impl Plugin for SelectPlugin { + fn build(&self, app: &mut App) { + app.register_type::() + } +} + + +#[derive(Default)] +pub struct SparsePlugin; + +impl Plugin for SparsePlugin { + fn build(&self, app: &mut App) { + app.register_type::(); + + app.add_systems(Update, select_sparse_handler); + } +} + + +impl KdPoint for Gaussian { + type Scalar = f32; + type Dim = U3; + + fn at(&self, i: usize) -> Self::Scalar { + self.position_visibility.position[i] + } +} + + +fn select_sparse_handler( + mut commands: Commands, + asset_server: Res, + gaussian_clouds_res: Res>, + mut selections: Query<( + Entity, + &Handle, + &mut SparseSelect, + )>, +) { + for ( + entity, + cloud_handle, + mut select, + ) in selections.iter_mut() { + if Some(LoadState::Loading) == asset_server.get_load_state(cloud_handle) { + continue; + } + + if Some(LoadState::Loading) == asset_server.get_load_state(cloud_handle) { + continue; + } + + if select.completed { + continue; + } + select.completed = true; + + let cloud = gaussian_clouds_res.get(cloud_handle).unwrap(); + let tree = KdTree::build_by_ordered_float(cloud.gaussian_iter().collect()); + + let new_selection = cloud.gaussian_iter() + .enumerate() + .filter(|(_idx, gaussian)| { + let neighbors = tree.within_radius(gaussian, select.radius); + + neighbors.len() < select.neighbor_threshold + }) + .map(|(idx, _gaussian)| idx) + .collect::() + .insert(new_selection); + } +} diff --git a/src/render/bindings.wgsl b/src/render/bindings.wgsl index 9f9f47c5..594d3710 100644 --- a/src/render/bindings.wgsl +++ b/src/render/bindings.wgsl @@ -11,14 +11,18 @@ struct GaussianUniforms { global_transform: mat4x4, global_scale: f32, count: u32, + count_root_ceil: u32, }; @group(1) @binding(0) var gaussian_uniforms: GaussianUniforms; + +// TODO: move these bindings to packed vs. planar +#ifdef PACKED_F32 struct Gaussian { @location(0) rotation: vec4, - @location(1) position: vec4, + @location(1) position_visibility: vec4, @location(2) scale_opacity: vec4, - sh: array, + sh: array, }; #ifdef READ_WRITE_POINTS @@ -26,6 +30,58 @@ struct Gaussian { #else @group(2) @binding(0) var points: array; #endif +#endif + + +#ifdef PLANAR_F32 +#ifdef READ_WRITE_POINTS +@group(2) @binding(0) var position_visibility: array>; +#else +@group(2) @binding(0) var position_visibility: array>; +#endif + +@group(2) @binding(1) var spherical_harmonics: array>; +@group(2) @binding(2) var rotation: array>; +@group(2) @binding(3) var scale_opacity: array>; +#endif + + +#ifdef PLANAR_F16 +#ifdef READ_WRITE_POINTS +@group(2) @binding(0) var position_visibility: array>; +#else +@group(2) @binding(0) var position_visibility: array>; +#endif + +@group(2) @binding(1) var spherical_harmonics: array>; +@group(2) @binding(2) var rotation_scale_opacity: array>; +#endif + + +#ifdef PLANAR_TEXTURE_F16 +@group(2) @binding(0) var position_visibility: texture_2d; + +#if SH_VEC4_PLANES == 1 +@group(2) @binding(1) var spherical_harmonics: texture_2d; +#else +@group(2) @binding(1) var spherical_harmonics: texture_2d_array; +#endif + +@group(2) @binding(2) var rotation_scale_opacity: texture_2d; +#endif + + +#ifdef PLANAR_TEXTURE_F32 +@group(2) @binding(0) var position_visibility: texture_2d; + +#if SH_VEC4_PLANES == 1 +@group(2) @binding(1) var spherical_harmonics: texture_2d; +#else +@group(2) @binding(1) var spherical_harmonics: texture_2d_array; +#endif + +@group(2) @binding(2) var rotation_scale_opacity: texture_2d; +#endif struct DrawIndirect { diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index 5061c968..175fd3de 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -2,7 +2,6 @@ view, globals, gaussian_uniforms, - points, sorting_pass_index, sorting, draw_indirect, @@ -10,18 +9,83 @@ output_entries, Entry, } -#import bevy_gaussian_splatting::color::{ +#import bevy_gaussian_splatting::depth::{ depth_to_rgb, } -#import bevy_gaussian_splatting::spherical_harmonics::spherical_harmonics_lookup #import bevy_gaussian_splatting::transform::{ world_to_clip, in_frustum, } +#ifdef PACKED +#import bevy_gaussian_splatting::packed::{ + get_position, + get_color, + get_rotation, + get_scale, + get_opacity, + get_visibility, +} +#endif + +#ifdef BUFFER_STORAGE +#import bevy_gaussian_splatting::planar::{ + get_position, + get_color, + get_rotation, + get_scale, + get_opacity, + get_visibility, +} +#endif +#ifdef BUFFER_TEXTURE +#import bevy_gaussian_splatting::texture::{ + get_position, + get_color, + get_rotation, + get_scale, + get_opacity, + get_visibility, + location, +} +#endif + + +#ifdef BUFFER_STORAGE @group(3) @binding(0) var sorted_entries: array; +fn get_entry(index: u32) -> Entry { + return sorted_entries[index]; +} +#endif + +#ifdef BUFFER_TEXTURE +@group(3) @binding(0) var sorted_entries: texture_2d; + +fn get_entry(index: u32) -> Entry { + let sample = textureLoad( + sorted_entries, + location(index), + 0, + ); + + return Entry( + sample.r, + sample.g, + ); +} +#endif + +#ifdef WEBGL2 +struct GaussianVertexOutput { + @builtin(position) position: vec4, + @location(0) color: vec4, + @location(1) conic: vec3, + @location(2) uv: vec2, + @location(3) major_minor: vec2, +}; +#else struct GaussianVertexOutput { @builtin(position) position: vec4, @location(0) @interpolate(flat) color: vec4, @@ -29,6 +93,7 @@ struct GaussianVertexOutput { @location(2) @interpolate(linear) uv: vec2, @location(3) @interpolate(linear) major_minor: vec2, }; +#endif // https://github.com/cvlab-epfl/gaussian-splatting-web/blob/905b3c0fb8961e42c79ef97e64609e82383ca1c2/src/shaders.ts#L185 @@ -207,19 +272,26 @@ fn vs_points( @builtin(vertex_index) vertex_index: u32, ) -> GaussianVertexOutput { var output: GaussianVertexOutput; - let splat_index = sorted_entries[instance_index][1]; - let discard_quad = sorted_entries[instance_index][0] == 0xFFFFFFFFu || splat_index == 0u; - if (discard_quad) { - output.color = vec4(0.0, 0.0, 0.0, 0.0); - output.position = vec4(0.0, 0.0, 0.0, 0.0); - return output; - } + let entry = get_entry(instance_index); + let splat_index = entry.value; + + var discard_quad = false; + + discard_quad |= entry.key == 0xFFFFFFFFu; // || splat_index == 0u; - let point = points[splat_index]; - let transformed_position = (gaussian_uniforms.global_transform * point.position).xyz; + let position = vec4(get_position(splat_index), 1.0); + + let transformed_position = (gaussian_uniforms.global_transform * position).xyz; let projected_position = world_to_clip(transformed_position); - if (!in_frustum(projected_position.xyz)) { + + discard_quad |= !in_frustum(projected_position.xyz); + +#ifdef DRAW_SELECTED + discard_quad |= get_visibility(splat_index) < 0.5; +#endif + + if (discard_quad) { output.color = vec4(0.0, 0.0, 0.0, 0.0); output.position = vec4(0.0, 0.0, 0.0, 0.0); return output; @@ -237,9 +309,14 @@ fn vs_points( let ray_direction = normalize(transformed_position - view.world_position); + var rgb = vec3(0.0); + #ifdef VISUALIZE_DEPTH - let min_position = (gaussian_uniforms.global_transform * points[sorted_entries[1][1]].position).xyz; - let max_position = (gaussian_uniforms.global_transform * points[sorted_entries[gaussian_uniforms.count - 1u][1]].position).xyz; + let first_position = vec4(get_position(get_entry(1u).value), 1.0); + let last_position = vec4(get_position(get_entry(gaussian_uniforms.count - 1u).value), 1.0); + + let min_position = (gaussian_uniforms.global_transform * first_position).xyz; + let max_position = (gaussian_uniforms.global_transform * last_position).xyz; let camera_position = view.world_position; @@ -247,24 +324,30 @@ fn vs_points( let max_distance = length(max_position - camera_position); let depth = length(transformed_position - camera_position); - let rgb = depth_to_rgb( + rgb = depth_to_rgb( depth, min_distance, max_distance, ); #else - let rgb = spherical_harmonics_lookup(ray_direction, point.sh); + rgb = get_color(splat_index, ray_direction); #endif + // TODO: precompute color, cov2d for every gaussian. cov2d only needs a single evaluation, while color needs to be evaluated every frame in SH degree > 0 mode output.color = vec4( rgb, - point.scale_opacity.a + get_opacity(splat_index), ); - // TODO: add depth color visualization +#ifdef HIGHLIGHT_SELECTED + if (get_visibility(splat_index) > 0.5) { + output.color = vec4(0.3, 1.0, 0.1, 1.0); + } +#endif - let cov2d = compute_cov2d(transformed_position, point.scale_opacity.rgb, point.rotation); + let cov2d = compute_cov2d(transformed_position, get_scale(splat_index), get_rotation(splat_index)); + // TODO: disable output.conic in obb mode let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y; let det_inv = 1.0 / det; let conic = vec3( diff --git a/src/render/mod.rs b/src/render/mod.rs index 4fb60409..8313da84 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -57,10 +57,16 @@ use bevy::{ use crate::{ gaussian::{ - Gaussian, - GaussianCloud, - GaussianCloudSettings, - MAX_SH_COEFF_COUNT, + cloud::GaussianCloud, + settings::{ + GaussianCloudDrawMode, + GaussianCloudSettings, + }, + }, + material::spherical_harmonics::{ + HALF_SH_COEFF_COUNT, + SH_COEFF_COUNT, + SH_VEC4_PLANES, }, morph::MorphPlugin, sort::{ @@ -69,11 +75,21 @@ use crate::{ }, }; +#[cfg(feature = "packed")] +mod packed; + +#[cfg(all(feature = "buffer_storage"))] +mod planar; + +#[cfg(feature = "buffer_texture")] +mod texture; + const BINDINGS_SHADER_HANDLE: Handle = Handle::weak_from_u128(675257236); -const COLOR_SHADER_HANDLE: Handle = Handle::weak_from_u128(51234253); const GAUSSIAN_SHADER_HANDLE: Handle = Handle::weak_from_u128(68294581); -const SPHERICAL_HARMONICS_SHADER_HANDLE: Handle = Handle::weak_from_u128(834667312); +const PACKED_SHADER_HANDLE: Handle = Handle::weak_from_u128(123623514); +const PLANAR_SHADER_HANDLE: Handle = Handle::weak_from_u128(72345231); +const TEXTURE_SHADER_HANDLE: Handle = Handle::weak_from_u128(26345735); const TRANSFORM_SHADER_HANDLE: Handle = Handle::weak_from_u128(734523534); @@ -91,22 +107,29 @@ impl Plugin for RenderPipelinePlugin { load_internal_asset!( app, - COLOR_SHADER_HANDLE, - "color.wgsl", + GAUSSIAN_SHADER_HANDLE, + "gaussian.wgsl", Shader::from_wgsl ); load_internal_asset!( app, - GAUSSIAN_SHADER_HANDLE, - "gaussian.wgsl", + PACKED_SHADER_HANDLE, + "packed.wgsl", + Shader::from_wgsl + ); + + load_internal_asset!( + app, + PLANAR_SHADER_HANDLE, + "planar.wgsl", Shader::from_wgsl ); load_internal_asset!( app, - SPHERICAL_HARMONICS_SHADER_HANDLE, - "spherical_harmonics.wgsl", + TEXTURE_SHADER_HANDLE, + "texture.wgsl", Shader::from_wgsl ); @@ -119,11 +142,15 @@ impl Plugin for RenderPipelinePlugin { app.add_plugins(RenderAssetPlugin::::default()); app.add_plugins(UniformComponentPlugin::::default()); + app.add_plugins(( MorphPlugin, SortPlugin, )); + #[cfg(feature = "buffer_texture")] + app.add_plugins(texture::BufferTexturePlugin); + if let Ok(render_app) = app.get_sub_app_mut(RenderApp) { render_app .add_render_command::() @@ -160,7 +187,11 @@ pub struct GpuGaussianSplattingBundle { #[derive(Debug, Clone)] pub struct GpuGaussianCloud { - pub gaussian_buffer: Buffer, + #[cfg(feature = "packed")] + pub packed: packed::PackedBuffers, + #[cfg(feature = "buffer_storage")] + pub planar: planar::PlanarBuffers, + pub count: usize, pub draw_indirect_buffer: Buffer, @@ -181,20 +212,7 @@ impl RenderAsset for GaussianCloud { gaussian_cloud: Self::ExtractedAsset, render_device: &mut SystemParamItem, ) -> Result> { - let gaussian_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor { - label: Some("gaussian cloud buffer"), - contents: bytemuck::cast_slice(gaussian_cloud.gaussians.as_slice()), - usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, - }); - - let count = gaussian_cloud.gaussians.len(); - - // let draw_indirect_buffer = render_device.create_buffer(&BufferDescriptor { - // label: Some("draw indirect buffer"), - // size: std::mem::size_of::() as u64, - // usage: BufferUsages::INDIRECT | BufferUsages::COPY_DST | BufferUsages::STORAGE | BufferUsages::COPY_SRC, - // mapped_at_creation: false, - // }); + let count = gaussian_cloud.len(); let draw_indirect_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor { label: Some("draw indirect buffer"), @@ -202,12 +220,19 @@ impl RenderAsset for GaussianCloud { usage: BufferUsages::INDIRECT | BufferUsages::COPY_DST | BufferUsages::STORAGE | BufferUsages::COPY_SRC, }); + // TODO: (extract GaussianCloud, TextureBuffers) when feature buffer_texture is enabled + Ok(GpuGaussianCloud { - gaussian_buffer, count, draw_indirect_buffer, + #[cfg(feature = "debug_gpu")] debug_gpu: gaussian_cloud, + + #[cfg(feature = "packed")] + packed: packed::prepare_cloud(render_device, &gaussian_cloud), + #[cfg(feature = "buffer_storage")] + planar: planar::prepare_cloud(render_device, &gaussian_cloud), }) } } @@ -222,15 +247,26 @@ fn queue_gaussians( pipeline_cache: Res, gaussian_clouds: Res>, sorted_entries: Res>, + mut views: Query<( + &ExtractedView, + &mut RenderPhase, + )>, + + #[cfg(feature = "buffer_storage")] gaussian_splatting_bundles: Query<( Entity, &Handle, &Handle, &GaussianCloudSettings, + (), )>, - mut views: Query<( - &ExtractedView, - &mut RenderPhase, + #[cfg(feature = "buffer_texture")] + gaussian_splatting_bundles: Query<( + Entity, + &Handle, + &Handle, + &GaussianCloudSettings, + &texture::GpuTextureBuffers, )>, ) { // TODO: condition this system based on GaussianCloudBindGroup attachment @@ -249,6 +285,7 @@ fn queue_gaussians( cloud_handle, sorted_entries_handle, settings, + _, ) in &gaussian_splatting_bundles { if gaussian_clouds.get(cloud_handle).is_none() { return; @@ -262,6 +299,7 @@ fn queue_gaussians( aabb: settings.aabb, visualize_bounding_box: settings.visualize_bounding_box, visualize_depth: settings.visualize_depth, + draw_mode: settings.draw_mode, }; let pipeline = pipelines.specialize(&pipeline_cache, &custom_pipeline, key); @@ -345,29 +383,20 @@ impl FromWorld for GaussianCloudPipeline { #[cfg(not(feature = "morph_particles"))] let read_only = true; - #[cfg(feature = "morph_particles")] let read_only = false; - let gaussian_cloud_layout = render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { - label: Some("gaussian_cloud_layout"), - entries: &[ - BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::all(), - ty: BindingType::Buffer { - ty: BufferBindingType::Storage { read_only }, - has_dynamic_offset: false, - min_binding_size: BufferSize::new(std::mem::size_of::() as u64), - }, - count: None, - }, - ], - }); + #[cfg(feature = "packed")] + let gaussian_cloud_layout = packed::get_bind_group_layout(&render_device, read_only); + #[cfg(feature = "buffer_storage")] + let gaussian_cloud_layout = planar::get_bind_group_layout(&render_device, read_only); + #[cfg(feature = "buffer_texture")] + let gaussian_cloud_layout = texture::get_bind_group_layout(&render_device, read_only); + #[cfg(feature = "buffer_storage")] let sorted_layout = render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { label: Some("sorted_layout"), - entries: &vec![ + entries: &[ BindGroupLayoutEntry { binding: 0, visibility: ShaderStages::VERTEX_FRAGMENT, @@ -380,6 +409,8 @@ impl FromWorld for GaussianCloudPipeline { }, ], }); + #[cfg(feature = "buffer_texture")] + let sorted_layout = texture::get_sorted_bind_group_layout(&render_device); GaussianCloudPipeline { gaussian_cloud_layout, @@ -450,13 +481,13 @@ impl Default for ShaderDefines { } pub fn shader_defs( - aabb: bool, - visualize_bounding_box: bool, - visualize_depth: bool, + key: GaussianCloudPipelineKey, ) -> Vec { let defines = ShaderDefines::default(); let mut shader_defs = vec![ - ShaderDefVal::UInt("MAX_SH_COEFF_COUNT".into(), MAX_SH_COEFF_COUNT as u32), + ShaderDefVal::UInt("SH_COEFF_COUNT".into(), SH_COEFF_COUNT as u32), + ShaderDefVal::UInt("HALF_SH_COEFF_COUNT".into(), HALF_SH_COEFF_COUNT as u32), + ShaderDefVal::UInt("SH_VEC4_PLANES".into(), SH_VEC4_PLANES as u32), ShaderDefVal::UInt("RADIX_BASE".into(), defines.radix_base), ShaderDefVal::UInt("RADIX_BITS_PER_DIGIT".into(), defines.radix_bits_per_digit), ShaderDefVal::UInt("RADIX_DIGIT_PLACES".into(), defines.radix_digit_places), @@ -469,44 +500,81 @@ pub fn shader_defs( ShaderDefVal::UInt("TEMPORAL_SORT_WINDOW_SIZE".into(), defines.temporal_sort_window_size), ]; - if aabb { + if key.aabb { shader_defs.push("USE_AABB".into()); } - if !aabb { + if !key.aabb { shader_defs.push("USE_OBB".into()); } - if visualize_bounding_box { + if key.visualize_bounding_box { shader_defs.push("VISUALIZE_BOUNDING_BOX".into()); } #[cfg(feature = "morph_particles")] shader_defs.push("READ_WRITE_POINTS".into()); - if visualize_depth { + if key.visualize_depth { shader_defs.push("VISUALIZE_DEPTH".into()); } + #[cfg(feature = "packed")] + shader_defs.push("PACKED".into()); + + #[cfg(feature = "buffer_storage")] + shader_defs.push("BUFFER_STORAGE".into()); + + #[cfg(feature = "buffer_texture")] + shader_defs.push("BUFFER_TEXTURE".into()); + + #[cfg(feature = "f16")] + shader_defs.push("F16".into()); + + #[cfg(feature = "f32")] + shader_defs.push("F32".into()); + + #[cfg(all(feature = "packed", feature = "f32"))] + shader_defs.push("PACKED_F32".into()); + + #[cfg(all(feature = "f16", feature = "buffer_storage"))] + shader_defs.push("PLANAR_F16".into()); + + #[cfg(all(feature = "f32", feature = "buffer_storage"))] + shader_defs.push("PLANAR_F32".into()); + + #[cfg(all(feature = "f16", feature = "buffer_texture"))] + shader_defs.push("PLANAR_TEXTURE_F16".into()); + + #[cfg(all(feature = "f32", feature = "buffer_texture"))] + shader_defs.push("PLANAR_TEXTURE_F32".into()); + + + #[cfg(feature = "webgl2")] + shader_defs.push("WEBGL2".into()); + + match key.draw_mode { + GaussianCloudDrawMode::All => {}, + GaussianCloudDrawMode::Selected => shader_defs.push("DRAW_SELECTED".into()), + GaussianCloudDrawMode::HighlightSelected => shader_defs.push("HIGHLIGHT_SELECTED".into()), + } + shader_defs } -#[derive(PartialEq, Eq, Hash, Clone, Copy)] +#[derive(PartialEq, Eq, Hash, Clone, Copy, Default)] pub struct GaussianCloudPipelineKey { pub aabb: bool, pub visualize_bounding_box: bool, pub visualize_depth: bool, + pub draw_mode: GaussianCloudDrawMode, } impl SpecializedRenderPipeline for GaussianCloudPipeline { type Key = GaussianCloudPipelineKey; fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor { - let shader_defs = shader_defs( - key.aabb, - key.visualize_bounding_box, - key.visualize_depth, - ); + let shader_defs = shader_defs(key); RenderPipelineDescriptor { label: Some("gaussian cloud render pipeline".into()), @@ -580,8 +648,10 @@ pub struct GaussianCloudUniform { pub transform: Mat4, pub global_scale: f32, pub count: u32, + pub count_root_ceil: u32, } +#[allow(clippy::type_complexity)] pub fn extract_gaussians( mut commands: Commands, mut prev_commands_len: Local, @@ -626,6 +696,7 @@ pub fn extract_gaussians( transform: settings.global_transform.compute_matrix(), global_scale: settings.global_scale, count: cloud.count as u32, + count_root_ceil: (cloud.count as f32).sqrt().ceil() as u32, }; commands_list.push(( entity, @@ -653,6 +724,7 @@ pub struct GaussianCloudBindGroup { pub sorted_bind_group: BindGroup, } +#[allow(clippy::too_many_arguments)] fn queue_gaussian_bind_group( mut commands: Commands, mut groups: ResMut, @@ -662,16 +734,29 @@ fn queue_gaussian_bind_group( asset_server: Res, gaussian_cloud_res: Res>, sorted_entries_res: Res>, + + #[cfg(feature = "buffer_storage")] + gaussian_clouds: Query<( + Entity, + &Handle, + &Handle, + )>, + #[cfg(feature = "buffer_texture")] gaussian_clouds: Query<( Entity, &Handle, &Handle, + &texture::GpuTextureBuffers, )>, + + #[cfg(feature = "buffer_texture")] + gpu_images: Res>, ) { let Some(model) = gaussian_uniforms.buffer() else { return; }; + // TODO: overloaded system, move to resource setup system groups.base_bind_group = Some(render_device.create_bind_group( "gaussian_uniform_bind_group", &gaussian_cloud_pipeline.gaussian_uniform_layout, @@ -687,11 +772,14 @@ fn queue_gaussian_bind_group( ], )); - for ( - entity, - cloud_handle, - sorted_entries_handle, - ) in gaussian_clouds.iter() { + for query in gaussian_clouds.iter() { + let entity = query.0; + let cloud_handle = query.1; + let sorted_entries_handle = query.2; + + #[cfg(feature = "buffer_texture")] + let texture_buffers = query.3; + // TODO: add asset loading indicator (and maybe streamed loading) if Some(LoadState::Loading) == asset_server.get_load_state(cloud_handle){ continue; @@ -709,38 +797,50 @@ fn queue_gaussian_bind_group( continue; } + #[cfg(not(feature = "buffer_texture"))] let cloud: &GpuGaussianCloud = gaussian_cloud_res.get(cloud_handle).unwrap(); + let sorted_entries = sorted_entries_res.get(sorted_entries_handle).unwrap(); + #[cfg(feature = "packed")] + let cloud_bind_group = packed::get_bind_group(&render_device, &gaussian_cloud_pipeline, &cloud); + #[cfg(feature = "buffer_storage")] + let cloud_bind_group = planar::get_bind_group(&render_device, &gaussian_cloud_pipeline, &cloud); + #[cfg(feature = "buffer_texture")] + let cloud_bind_group = texture_buffers.bind_group.clone(); + + #[cfg(feature = "buffer_storage")] + let sorted_bind_group = render_device.create_bind_group( + "render_sorted_bind_group", + &gaussian_cloud_pipeline.sorted_layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::Buffer(BufferBinding { + buffer: &sorted_entries.sorted_entry_buffer, + offset: 0, + size: BufferSize::new((cloud.count * std::mem::size_of::<(u32, u32)>()) as u64), + }), + }, + ], + ); + #[cfg(feature = "buffer_texture")] + let sorted_bind_group = render_device.create_bind_group( + Some("render_sorted_bind_group"), + &gaussian_cloud_pipeline.sorted_layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::TextureView( + &gpu_images.get(&sorted_entries.texture).unwrap().texture_view + ), + }, + ], + ); + commands.entity(entity).insert(GaussianCloudBindGroup { - cloud_bind_group: render_device.create_bind_group( - "gaussian_cloud_bind_group", - &gaussian_cloud_pipeline.gaussian_cloud_layout, - &[ - BindGroupEntry { - binding: 0, - resource: BindingResource::Buffer(BufferBinding { - buffer: &cloud.gaussian_buffer, - offset: 0, - size: BufferSize::new(cloud.gaussian_buffer.size()), - }), - }, - ], - ), - sorted_bind_group: render_device.create_bind_group( - "render_sorted_bind_group", - &gaussian_cloud_pipeline.sorted_layout, - &[ - BindGroupEntry { - binding: 0, - resource: BindingResource::Buffer(BufferBinding { - buffer: &sorted_entries.sorted_entry_buffer, - offset: 0, - size: BufferSize::new((cloud.count as usize * std::mem::size_of::<(u32, u32)>()) as u64), - }), - }, - ], - ), + cloud_bind_group, + sorted_bind_group, }); } } @@ -885,9 +985,13 @@ impl RenderCommand

for DrawGaussianInstanced { None => return RenderCommandResult::Failure, }; - pass.set_bind_group(2, &bind_groups.cloud_bind_group, &[]); + pass.set_bind_group(2, &bind_groups.cloud_bind_group, &[]); // TODO: abstract source of cloud_bind_group (e.g. packed vs. planar) pass.set_bind_group(3, &bind_groups.sorted_bind_group, &[]); + #[cfg(feature = "webgl2")] + pass.draw(0..4, 0..gpu_gaussian_cloud.count as u32); + + #[cfg(not(feature = "webgl2"))] pass.draw_indirect(&gpu_gaussian_cloud.draw_indirect_buffer, 0); RenderCommandResult::Success diff --git a/src/render/packed.rs b/src/render/packed.rs new file mode 100644 index 00000000..99e8d692 --- /dev/null +++ b/src/render/packed.rs @@ -0,0 +1,90 @@ +use bevy::render::{ + render_resource::{ + Buffer, + BufferInitDescriptor, + BufferUsages, + Extent3d, + ShaderType, + TextureDimension, + TextureFormat, + }, + renderer::RenderDevice, +}; + +use crate::{ + gaussian::{ + cloud::GaussianCloud, + packed::Gaussian, + }, + render::{ + GaussianCloudPipeline, + GpuGaussianCloud, + }, +}; + + +#[derive(Debug, Clone)] +pub struct PackedBuffers { + gaussians: Buffer, +} + + +pub fn prepare_cloud( + render_device: &RenderDevice, + cloud: &GaussianCloud, +) -> PackedBuffers { + let gaussians = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("packed_gaussian_cloud_buffer"), + contents: bytemuck::cast_slice(cloud.gaussian_iter().collect::>().as_slice()), + usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, + }); + + PackedBuffers { + gaussians, + } +} + + +pub fn get_bind_group_layout( + render_device: &RenderDevice, + read_only: bool +) -> BindGroupLayout { + render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: Some("packed_gaussian_cloud_layout"), + entries: &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + ], + }) +} + + +#[cfg(feature = "packed")] +pub fn get_bind_group( + render_device: &RenderDevice, + gaussian_cloud_pipeline: &GaussianCloudPipeline, + cloud: &GpuGaussianCloud, +) -> BindGroup { + render_device.create_bind_group( + "packed_gaussian_cloud_bind_group", + &gaussian_cloud_pipeline.gaussian_cloud_layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::Buffer(BufferBinding { + buffer: &cloud.packed.gaussians, + offset: 0, + size: BufferSize::new(cloud.packed.gaussians.size()), + }), + }, + ], + ) +} diff --git a/src/render/packed.wgsl b/src/render/packed.wgsl new file mode 100644 index 00000000..cfe3ba02 --- /dev/null +++ b/src/render/packed.wgsl @@ -0,0 +1,45 @@ +#define_import_path bevy_gaussian_splatting::packed + +#import bevy_gaussian_splatting::bindings::points +#import bevy_gaussian_splatting::spherical_harmonics::{ + spherical_harmonics_lookup, + srgb_to_linear, +} + + +#ifdef PACKED_F32 + +fn get_position(index: u32) -> vec3 { + return points[index].position_visibility.xyz; +} + +fn get_color( + index: u32, + ray_direction: vec3, +) -> vec3 { + let sh = get_spherical_harmonics(index); + let color = spherical_harmonics_lookup(ray_direction, sh); + return srgb_to_linear(color); +} + +fn get_spherical_harmonics(index: u32) -> array { + return points[index].sh; +} + +fn get_rotation(index: u32) -> vec4 { + return points[index].rotation; +} + +fn get_scale(index: u32) -> vec3 { + return points[index].scale_opacity.xyz; +} + +fn get_opacity(index: u32) -> f32 { + return points[index].scale_opacity.w; +} + +fn get_visibility(index: u32) -> f32 { + return points[index].position_visibility.w; +} + +#endif diff --git a/src/render/planar.rs b/src/render/planar.rs new file mode 100644 index 00000000..95a66ee8 --- /dev/null +++ b/src/render/planar.rs @@ -0,0 +1,290 @@ +#[allow(unused_imports)] +use bevy::render::{ + render_resource::*, + renderer::RenderDevice, +}; + +#[allow(unused_imports)] +use crate::{ + gaussian::{ + cloud::GaussianCloud, + f32::{ + PositionVisibility, + Rotation, + ScaleOpacity, + }, + }, + render::{ + GaussianCloudPipeline, + GpuGaussianCloud, + }, + material::spherical_harmonics::SphericalHarmonicCoefficients, +}; + +#[cfg(feature = "f16")] +use crate::gaussian::f16::RotationScaleOpacityPacked128; + +#[derive(Debug, Clone)] +pub struct PlanarBuffers { + position_visibility: Buffer, + spherical_harmonics: Buffer, + + #[cfg(feature = "f16")] + rotation_scale_opacity: Buffer, + + #[cfg(feature = "f32")] + rotation: Buffer, + #[cfg(feature = "f32")] + scale_opacity: Buffer, +} + + +#[cfg(feature = "f16")] +pub fn prepare_cloud( + render_device: &RenderDevice, + cloud: &GaussianCloud, +) -> PlanarBuffers { + let position_visibility = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("planar_position_visibility_buffer"), + contents: bytemuck::cast_slice(cloud.position_visibility.as_slice()), + usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, + }); + + let rotation_scale_opacity = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("planar_rotation_scale_opacity_buffer"), + contents: bytemuck::cast_slice(cloud.rotation_scale_opacity_packed128.as_slice()), + usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, + }); + + let spherical_harmonics = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("planar_spherical_harmonics_buffer"), + contents: bytemuck::cast_slice(cloud.spherical_harmonic.as_slice()), + usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, + }); + + PlanarBuffers { + position_visibility, + rotation_scale_opacity, + spherical_harmonics, + } +} + + +#[cfg(feature = "f32")] +pub fn prepare_cloud( + render_device: &RenderDevice, + cloud: &GaussianCloud, +) -> PlanarBuffers { + let position_visibility = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("planar_f32_position_visibility_buffer"), + contents: bytemuck::cast_slice(cloud.position_visibility.as_slice()), + usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, + }); + + let rotation = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("planar_f32_rotation_buffer"), + contents: bytemuck::cast_slice(cloud.rotation.as_slice()), + usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, + }); + + let scale_opacity = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("planar_f32_scale_opacity_buffer"), + contents: bytemuck::cast_slice(cloud.scale_opacity.as_slice()), + usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, + }); + + let spherical_harmonics = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("planar_f32_spherical_harmonics_buffer"), + contents: bytemuck::cast_slice(cloud.spherical_harmonic.as_slice()), + usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, + }); + + PlanarBuffers { + position_visibility, + rotation, + scale_opacity, + spherical_harmonics, + } +} + + +#[cfg(feature = "f16")] +pub fn get_bind_group_layout( + render_device: &RenderDevice, + read_only: bool +) -> BindGroupLayout { + render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: Some("planar_f16_gaussian_cloud_layout"), + entries: &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + ], + }) +} + + +#[cfg(feature = "f32")] +pub fn get_bind_group_layout( + render_device: &RenderDevice, + read_only: bool +) -> BindGroupLayout { + render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: Some("planar_f32_gaussian_cloud_layout"), + entries: &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 3, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + ], + }) +} + + +#[cfg(all(feature = "planar", feature = "f16"))] +pub fn get_bind_group( + render_device: &RenderDevice, + gaussian_cloud_pipeline: &GaussianCloudPipeline, + cloud: &GpuGaussianCloud, +) -> BindGroup { + render_device.create_bind_group( + "planar_gaussian_cloud_bind_group", + &gaussian_cloud_pipeline.gaussian_cloud_layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::Buffer(BufferBinding { + buffer: &cloud.planar.position_visibility, + offset: 0, + size: BufferSize::new(cloud.planar.position_visibility.size()), + }), + }, + BindGroupEntry { + binding: 1, + resource: BindingResource::Buffer(BufferBinding { + buffer: &cloud.planar.spherical_harmonics, + offset: 0, + size: BufferSize::new(cloud.planar.spherical_harmonics.size()), + }), + }, + BindGroupEntry { + binding: 2, + resource: BindingResource::Buffer(BufferBinding { + buffer: &cloud.planar.rotation_scale_opacity, + offset: 0, + size: BufferSize::new(cloud.planar.rotation_scale_opacity.size()), + }), + }, + ], + ) +} + + +#[cfg(all(feature = "planar", feature = "f32"))] +pub fn get_bind_group( + render_device: &RenderDevice, + gaussian_cloud_pipeline: &GaussianCloudPipeline, + cloud: &GpuGaussianCloud, +) -> BindGroup { + render_device.create_bind_group( + "planar_gaussian_cloud_bind_group", + &gaussian_cloud_pipeline.gaussian_cloud_layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::Buffer(BufferBinding { + buffer: &cloud.planar.position_visibility, + offset: 0, + size: BufferSize::new(cloud.planar.position_visibility.size()), + }), + }, + BindGroupEntry { + binding: 1, + resource: BindingResource::Buffer(BufferBinding { + buffer: &cloud.planar.spherical_harmonics, + offset: 0, + size: BufferSize::new(cloud.planar.spherical_harmonics.size()), + }), + }, + BindGroupEntry { + binding: 2, + resource: BindingResource::Buffer(BufferBinding { + buffer: &cloud.planar.rotation, + offset: 0, + size: BufferSize::new(cloud.planar.rotation.size()), + }), + }, + BindGroupEntry { + binding: 3, + resource: BindingResource::Buffer(BufferBinding { + buffer: &cloud.planar.scale_opacity, + offset: 0, + size: BufferSize::new(cloud.planar.scale_opacity.size()), + }), + }, + ], + ) +} diff --git a/src/render/planar.wgsl b/src/render/planar.wgsl new file mode 100644 index 00000000..b1e85d11 --- /dev/null +++ b/src/render/planar.wgsl @@ -0,0 +1,99 @@ +#define_import_path bevy_gaussian_splatting::planar + +#import bevy_gaussian_splatting::bindings::{ + position_visibility, + spherical_harmonics, + rotation, + rotation_scale_opacity, + scale_opacity, +}; +#import bevy_gaussian_splatting::spherical_harmonics::{ + spherical_harmonics_lookup, + srgb_to_linear, +} + + +fn get_color( + index: u32, + ray_direction: vec3, +) -> vec3 { + let sh = get_spherical_harmonics(index); + let color = spherical_harmonics_lookup(ray_direction, sh); + return srgb_to_linear(color); +} + + +#ifdef PLANAR_F16 + +fn get_position(index: u32) -> vec3 { + return position_visibility[index].xyz; +} + +fn get_spherical_harmonics(index: u32) -> array { + var coefficients: array; + + for (var i = 0u; i < #{HALF_SH_COEFF_COUNT}u; i = i + 1u) { + let values = unpack2x16float(spherical_harmonics[index][i]); + + coefficients[i * 2u] = values[0]; + coefficients[i * 2u + 1u] = values[1]; + } + + return coefficients; +} + +fn get_rotation(index: u32) -> vec4 { + let q0 = unpack2x16float(rotation_scale_opacity[index].x); + let q1 = unpack2x16float(rotation_scale_opacity[index].y); + + return vec4( + q0.yx, + q1.yx, + ); +} + +fn get_scale(index: u32) -> vec3 { + let s0 = unpack2x16float(rotation_scale_opacity[index].z); + let s1 = unpack2x16float(rotation_scale_opacity[index].w); + + return vec3( + s0.yx, + s1.y, + ); +} + +fn get_opacity(index: u32) -> f32 { + return unpack2x16float(rotation_scale_opacity[index].w).x; +} + +fn get_visibility(index: u32) -> f32 { + return position_visibility[index].w; +} +#endif + + +#ifdef PLANAR_F32 +fn get_position(index: u32) -> vec3 { + return position_visibility[index].xyz; +} + +fn get_spherical_harmonics(index: u32) -> array { + return spherical_harmonics[index]; +} + +fn get_rotation(index: u32) -> vec4 { + return rotation[index]; +} + +fn get_scale(index: u32) -> vec3 { + return scale_opacity[index].xyz; +} + +fn get_opacity(index: u32) -> f32 { + return scale_opacity[index].w; +} + +fn get_visibility(index: u32) -> f32 { + return position_visibility[index].w; +} +#endif diff --git a/src/render/texture.rs b/src/render/texture.rs new file mode 100644 index 00000000..430ae58f --- /dev/null +++ b/src/render/texture.rs @@ -0,0 +1,420 @@ +use bevy::{ + prelude::*, + asset::LoadState, + ecs::query::QueryItem, + render::{ + extract_component::{ + ExtractComponent, + ExtractComponentPlugin, + }, + Render, + RenderApp, + RenderSet, + render_asset::RenderAssets, + render_resource::{ + BindGroup, + BindGroupLayout, + BindGroupLayoutDescriptor, + BindGroupLayoutEntry, + BindGroupEntry, + BindingType, + BindingResource, + Extent3d, + TextureDimension, + TextureFormat, + TextureSampleType, + TextureUsages, + TextureViewDimension, + ShaderStages, + }, + renderer::RenderDevice, + }, +}; +use static_assertions::assert_cfg; + +#[allow(unused_imports)] +use crate::{ + gaussian::{ + cloud::GaussianCloud, + f32::{ + PositionVisibility, + Rotation, + ScaleOpacity, + }, + }, + material::spherical_harmonics::{ + SH_COEFF_COUNT, + SH_VEC4_PLANES, + SphericalHarmonicCoefficients, + }, + render::{ + GaussianCloudPipeline, + GpuGaussianCloud, + }, +}; + + +// TODO: support loading from directory of images + + +assert_cfg!( + feature = "planar", + "texture rendering is only supported with the `planar` feature enabled", +); + +assert_cfg!( + not(feature = "f32"), + "f32 texture support is not implemented yet", +); + + +#[derive(Component, Clone, Debug, Reflect)] +pub struct TextureBuffers { + position_visibility: Handle, + spherical_harmonics: Handle, + + #[cfg(feature = "f16")] + rotation_scale_opacity: Handle, + + #[cfg(feature = "f32")] + rotation: Handle, + #[cfg(feature = "f32")] + scale_opacity: Handle, +} + +impl ExtractComponent for TextureBuffers { + type Query = &'static Self; + + type Filter = (); + type Out = Self; + + fn extract_component(texture_buffers: QueryItem<'_, Self::Query>) -> Option { + texture_buffers.clone().into() + } +} + + +#[derive(Default)] +pub struct BufferTexturePlugin; + +impl Plugin for BufferTexturePlugin { + fn build(&self, app: &mut App) { + app.register_type::(); + app.add_plugins(ExtractComponentPlugin::::default()); + + app.add_systems(Update, queue_textures); + + let render_app = app.sub_app_mut(RenderApp); + render_app.add_systems( + Render, + queue_gpu_texture_buffers.in_set(RenderSet::PrepareAssets), + ); + } +} + + +#[derive(Component, Clone, Debug)] +pub struct GpuTextureBuffers { + pub bind_group: BindGroup, +} + +pub fn queue_gpu_texture_buffers( + mut commands: Commands, + // gaussian_cloud_pipeline: Res, + pipeline: ResMut, + render_device: ResMut, + gpu_images: Res>, + clouds: Query<( + Entity, + &TextureBuffers, + )>, +) { + for (entity, texture_buffers,) in clouds.iter() { + #[cfg(feature = "f16")] + let bind_group = render_device.create_bind_group( + Some("texture_gaussian_cloud_bind_group"), + &pipeline.gaussian_cloud_layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::TextureView( + &gpu_images.get(&texture_buffers.position_visibility).unwrap().texture_view + ), + }, + BindGroupEntry { + binding: 1, + resource: BindingResource::TextureView( + &gpu_images.get(&texture_buffers.spherical_harmonics).unwrap().texture_view + ), + }, + BindGroupEntry { + binding: 2, + resource: BindingResource::TextureView( + &gpu_images.get(&texture_buffers.rotation_scale_opacity).unwrap().texture_view + ), + }, + ], + ); + + #[cfg(feature = "f32")] + let bind_group = render_device.create_bind_group( + Some("texture_gaussian_cloud_bind_group"), + &pipeline.gaussian_cloud_layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::TextureView( + &gpu_images.get(&texture_buffers.position_visibility).unwrap().texture_view + ), + }, + BindGroupEntry { + binding: 1, + resource: BindingResource::TextureView( + &gpu_images.get(&texture_buffers.spherical_harmonics).unwrap().texture_view + ), + }, + BindGroupEntry { + binding: 2, + resource: BindingResource::TextureView( + &gpu_images.get(&texture_buffers.rotation).unwrap().texture_view + ), + }, + BindGroupEntry { + binding: 3, + resource: BindingResource::TextureView( + &gpu_images.get(&texture_buffers.scale_opacity).unwrap().texture_view + ), + }, + ], + ); + + commands.entity(entity).insert(GpuTextureBuffers { bind_group }); + } +} + + +// TODO: support asset change detection and reupload +fn queue_textures( + mut commands: Commands, + asset_server: Res, + gaussian_cloud_res: Res>, + mut images: ResMut>, + clouds: Query<( + Entity, + &Handle, + Without, + )>, +) { + for (entity, cloud_handle, _) in clouds.iter() { + if Some(LoadState::Loading) == asset_server.get_load_state(cloud_handle){ + continue; + } + + if gaussian_cloud_res.get(cloud_handle).is_none() { + continue; + } + + let cloud = gaussian_cloud_res.get(cloud_handle).unwrap(); + + let square = cloud.len_sqrt_ceil() as u32; + let extent_1d = Extent3d { + width: square, + height: square, // TODO: shrink height to save memory (consider fixed width) + depth_or_array_layers: 1, + }; + + let mut position_visibility = Image::new( + extent_1d, + TextureDimension::D2, + bytemuck::cast_slice(cloud.position_visibility.as_slice()).to_vec(), + TextureFormat::Rgba32Float, + ); + position_visibility.texture_descriptor.usage = TextureUsages::COPY_DST | TextureUsages::TEXTURE_BINDING; + let position_visibility = images.add(position_visibility); + + let texture_buffers: TextureBuffers; + + #[cfg(feature = "f16")] + { + let planar_spherical_harmonics: Vec = (0..SH_VEC4_PLANES) + .flat_map(|plane_index| { + cloud.spherical_harmonic.iter() + .flat_map(move |sh| { + let start_index = plane_index * 4; + let end_index = std::cmp::min(start_index + 4, sh.coefficients.len()); + + let mut depthwise = sh.coefficients[start_index..end_index].to_vec(); + depthwise.resize(4, 0); + + depthwise + }) + }) + .collect(); + + let mut spherical_harmonics = Image::new( + Extent3d { + width: square, + height: square, + depth_or_array_layers: SH_VEC4_PLANES as u32, + }, + TextureDimension::D2, + bytemuck::cast_slice(planar_spherical_harmonics.as_slice()).to_vec(), + TextureFormat::Rgba32Uint, + ); + spherical_harmonics.texture_descriptor.usage = TextureUsages::COPY_DST | TextureUsages::TEXTURE_BINDING; + let spherical_harmonics = images.add(spherical_harmonics); + + let mut rotation_scale_opacity = Image::new( + extent_1d, + TextureDimension::D2, + bytemuck::cast_slice(cloud.rotation_scale_opacity_packed128.as_slice()).to_vec(), + TextureFormat::Rgba32Uint, + ); + rotation_scale_opacity.texture_descriptor.usage = TextureUsages::COPY_DST | TextureUsages::TEXTURE_BINDING; + let rotation_scale_opacity = images.add(rotation_scale_opacity); + + texture_buffers = TextureBuffers { + position_visibility, + spherical_harmonics, + rotation_scale_opacity, + }; + } + + #[cfg(feature = "f32")] + { + texture_buffers = TextureBuffers { + position_visibility, + spherical_harmonics: todo!(), + rotation: todo!(), + scale_opacity: todo!(), + }; + } + + commands.entity(entity).insert(texture_buffers); + } +} + + +pub fn get_sorted_bind_group_layout( + render_device: &RenderDevice, +) -> BindGroupLayout { + render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: Some("texture_sorted_layout"), + entries: &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::all(), + ty: BindingType::Texture { + view_dimension: TextureViewDimension::D2, + sample_type: TextureSampleType::Uint, + multisampled: false, + }, + count: None, + }, + ], + }) +} + + +#[cfg(feature = "f16")] +pub fn get_bind_group_layout( + render_device: &RenderDevice, + _read_only: bool +) -> BindGroupLayout { + let sh_view_dimension = if SH_VEC4_PLANES == 1 { + TextureViewDimension::D2 + } else { + TextureViewDimension::D2Array + }; + + render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: Some("texture_f16_gaussian_cloud_layout"), + entries: &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::all(), + ty: BindingType::Texture { + view_dimension: TextureViewDimension::D2, + sample_type: TextureSampleType::Float { + filterable: false, + }, + multisampled: false, + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::all(), + ty: BindingType::Texture { + view_dimension: sh_view_dimension, + sample_type: TextureSampleType::Uint, + multisampled: false, + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::all(), + ty: BindingType::Texture { + view_dimension: TextureViewDimension::D2, + sample_type: TextureSampleType::Uint, + multisampled: false, + }, + count: None, + }, + ], + }) +} + + +#[cfg(feature = "f32")] +pub fn get_bind_group_layout( + render_device: &RenderDevice, + read_only: bool +) -> BindGroupLayout { + render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: Some("texture_f32_gaussian_cloud_layout"), + entries: &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 3, + visibility: ShaderStages::all(), + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + ], + }) +} diff --git a/src/render/texture.wgsl b/src/render/texture.wgsl new file mode 100644 index 00000000..1950b77d --- /dev/null +++ b/src/render/texture.wgsl @@ -0,0 +1,342 @@ +#define_import_path bevy_gaussian_splatting::texture + +#import bevy_gaussian_splatting::bindings::{ + gaussian_uniforms, + position_visibility, + spherical_harmonics, + rotation, + rotation_scale_opacity, + scale_opacity, +}; +#import bevy_gaussian_splatting::spherical_harmonics::{ + shc, + srgb_to_linear, +} + + +fn location(index: u32) -> vec2 { + return vec2( + i32(index) % i32(gaussian_uniforms.count_root_ceil), + i32(index) / i32(gaussian_uniforms.count_root_ceil), + ); +} + + +#ifdef PLANAR_TEXTURE_F16 + +fn get_position(index: u32) -> vec3 { + let sample = textureLoad( + position_visibility, + location(index), + 0, + ); + + return sample.xyz; +} + +fn get_sh_vec( + index: u32, + plane: i32, +) -> vec4 { +#if SH_VEC4_PLANES == 1 + return textureLoad( + spherical_harmonics, + location(index), + plane, + ); +#else + return textureLoad( + spherical_harmonics, + location(index), + plane, + 0, + ); +#endif +} + +#ifdef WEBGL2 +fn get_color( + index: u32, + ray_direction: vec3, +) -> vec3 { + let s0 = get_sh_vec(index, 0); + + let v0 = unpack2x16float(s0.x); + let v1 = unpack2x16float(s0.y); + let v2 = unpack2x16float(s0.z); + let v3 = unpack2x16float(s0.w); + + let rds = ray_direction * ray_direction; + var color = vec3(0.5); + + color += shc[ 0] * vec3( + v0.x, + v0.y, + v1.x, + ); + +#if SH_COEFF_COUNT > 11 + let r1 = vec3( + v1.y, + v2.x, + v2.y, + ); + + let s1 = get_sh_vec(index, 1); + let v4 = unpack2x16float(s1.x); + let v5 = unpack2x16float(s1.y); + let v6 = unpack2x16float(s1.z); + let v7 = unpack2x16float(s1.w); + + let r2 = vec3( + v3.x, + v3.y, + v4.x, + ); + + let r3 = vec3( + v4.y, + v5.x, + v5.y, + ); + + color += shc[ 1] * r1 * ray_direction.y; + color += shc[ 2] * r2 * ray_direction.z; + color += shc[ 3] * r3 * ray_direction.x; +#endif + +#if SH_COEFF_COUNT > 26 + let r4 = vec3( + v6.x, + v6.y, + v7.x, + ); + + let s2 = get_sh_vec(index, 2); + let v8 = unpack2x16float(s2.x); + let v9 = unpack2x16float(s2.y); + let v10 = unpack2x16float(s2.z); + let v11 = unpack2x16float(s2.w); + + let r5 = vec3( + v7.y, + v8.x, + v8.y, + ); + + let r6 = vec3( + v9.x, + v9.y, + v10.x, + ); + + let r7 = vec3( + v10.y, + v11.x, + v11.y, + ); + + let s3 = get_sh_vec(index, 3); + let v12 = unpack2x16float(s3.x); + let v13 = unpack2x16float(s3.y); + let v14 = unpack2x16float(s3.z); + let v15 = unpack2x16float(s3.w); + + let r8 = vec3( + v12.x, + v12.y, + v13.x, + ); + + color += shc[ 4] * r4 * ray_direction.x * ray_direction.y; + color += shc[ 5] * r5 * ray_direction.y * ray_direction.z; + color += shc[ 6] * r6 * (2.0 * rds.z - rds.x - rds.y); + color += shc[ 7] * r7 * ray_direction.x * ray_direction.z; + color += shc[ 8] * r8 * (rds.x - rds.y); +#endif + +#if SH_COEFF_COUNT > 47 + let r9 = vec3( + v13.y, + v14.x, + v14.y, + ); + + let s4 = get_sh_vec(index, 4); + let v16 = unpack2x16float(s4.x); + let v17 = unpack2x16float(s4.y); + let v18 = unpack2x16float(s4.z); + let v19 = unpack2x16float(s4.w); + + let r10 = vec3( + v15.x, + v15.y, + v16.x, + ); + + let r11 = vec3( + v16.y, + v17.x, + v17.y, + ); + + let r12 = vec3( + v18.x, + v18.y, + v19.x, + ); + + let s5 = get_sh_vec(index, 5); + let v20 = unpack2x16float(s5.x); + let v21 = unpack2x16float(s5.y); + let v22 = unpack2x16float(s5.z); + let v23 = unpack2x16float(s5.w); + + let r13 = vec3( + v19.y, + v20.x, + v20.y, + ); + + let r14 = vec3( + v21.x, + v21.y, + v22.x, + ); + + let r15 = vec3( + v22.y, + v23.x, + v23.y, + ); + + color += shc[ 9] * r9 * ray_direction.y * (3.0 * rds.x - rds.y); + color += shc[10] * r10 * ray_direction.x * ray_direction.y * ray_direction.z; + color += shc[11] * r11 * ray_direction.y * (4.0 * rds.z - rds.x - rds.y); + color += shc[12] * r12 * ray_direction.z * (2.0 * rds.z - 3.0 * rds.x - 3.0 * rds.y); + color += shc[13] * r13 * ray_direction.x * (4.0 * rds.z - rds.x - rds.y); + color += shc[14] * r14 * ray_direction.z * (rds.x - rds.y); + color += shc[15] * r15 * ray_direction.x * (rds.x - 3.0 * rds.y); +#endif + + return srgb_to_linear(color); +} +#else +fn get_spherical_harmonics(index: u32) -> array { + var coefficients: array; + + for (var i = 0u; i < #{SH_VEC4_PLANES}u; i = i + 1u) { + let sample = get_sh_vec(index, i32(i)); + + let v0 = unpack2x16float(sample.x); + let v1 = unpack2x16float(sample.y); + let v2 = unpack2x16float(sample.z); + let v3 = unpack2x16float(sample.w); + + let base_index = i * 8u; + coefficients[base_index ] = v0.x; + coefficients[base_index + 1u] = v0.y; + + coefficients[base_index + 2u] = v1.x; + coefficients[base_index + 3u] = v1.y; + + coefficients[base_index + 4u] = v2.x; + coefficients[base_index + 5u] = v2.y; + + coefficients[base_index + 6u] = v3.x; + coefficients[base_index + 7u] = v3.y; + } + + return coefficients; +} + +fn get_color( + index: u32, + ray_direction: vec3, +) -> vec3 { + let sh = get_spherical_harmonics(index); + let color = spherical_harmonics_lookup(ray_direction, sh); + return srgb_to_linear(color); +} +#endif + +fn get_rotation(index: u32) -> vec4 { + let sample = textureLoad( + rotation_scale_opacity, + location(index), + 0, + ); + + let q0 = unpack2x16float(sample.x); + let q1 = unpack2x16float(sample.y); + + return vec4( + q0.yx, + q1.yx, + ); +} + +fn get_scale(index: u32) -> vec3 { + let sample = textureLoad( + rotation_scale_opacity, + location(index), + 0, + ); + + let s0 = unpack2x16float(sample.z); + let s1 = unpack2x16float(sample.w); + + return vec3( + s0.yx, + s1.y, + ); +} + +fn get_opacity(index: u32) -> f32 { + let sample = textureLoad( + rotation_scale_opacity, + location(index), + 0, + ); + + return unpack2x16float(sample.w).x; +} + +fn get_visibility(index: u32) -> f32 { + let sample = textureLoad( + position_visibility, + location(index), + 0, + ); + + return sample.w; +} +#endif + + +// TODO: support f32 +#ifdef PLANAR_TEXTURE_F32 +fn get_position(index: u32) -> vec3 { + return position_visibility[index].xyz; +} + +fn get_spherical_harmonics(index: u32) -> array { + return spherical_harmonics[index]; +} + +fn get_rotation(index: u32) -> vec4 { + return rotation[index]; +} + +fn get_scale(index: u32) -> vec3 { + return scale_opacity[index].xyz; +} + +fn get_opacity(index: u32) -> f32 { + return scale_opacity[index].w; +} + +fn get_visibility(index: u32) -> f32 { + return position_visibility[index].w; +} +#endif diff --git a/src/sort/mod.rs b/src/sort/mod.rs index f26e1168..e44baf1f 100644 --- a/src/sort/mod.rs +++ b/src/sort/mod.rs @@ -7,12 +7,12 @@ use bevy::{ }, reflect::TypeUuid, render::{ + render_resource::*, render_asset::{ RenderAsset, RenderAssetPlugin, PrepareAssetError, }, - render_resource::*, renderer::RenderDevice, }, }; @@ -35,7 +35,7 @@ pub mod radix; pub mod rayon; #[cfg(feature = "sort_std")] -pub mod std; +pub mod std; // rename to std_sort.rs to avoid name conflict with std crate assert_cfg!( @@ -107,10 +107,36 @@ impl Plugin for SortPlugin { app.add_plugins(RenderAssetPlugin::::default()); app.add_systems(Update, auto_insert_sorted_entries); + + #[cfg(feature = "buffer_texture")] + app.add_systems(PostUpdate, update_textures_on_change); + } +} + + +#[cfg(feature = "buffer_texture")] +fn update_textures_on_change( + mut images: ResMut>, + mut ev_asset: EventReader>, + sorted_entries_res: Res>, +) { + for ev in ev_asset.read() { + match ev { + AssetEvent::Modified { id } => { + let sorted_entries = sorted_entries_res.get(*id).unwrap(); + let image = images.get_mut(&sorted_entries.texture).unwrap(); + + image.data = bytemuck::cast_slice(sorted_entries.sorted.as_slice()).to_vec(); + }, + AssetEvent::Added { id: _ } => {}, + AssetEvent::Removed { id: _ } => {}, + AssetEvent::LoadedWithDependencies { id: _ } => {}, + } } } +#[allow(clippy::type_complexity)] fn auto_insert_sorted_entries( mut commands: Commands, asset_server: Res, @@ -122,6 +148,9 @@ fn auto_insert_sorted_entries( &GaussianCloudSettings, Without>, )>, + + #[cfg(feature = "buffer_texture")] + mut images: ResMut>, ) { for ( entity, @@ -144,16 +173,34 @@ fn auto_insert_sorted_entries( } let cloud = cloud.unwrap(); + let sorted: Vec = (0..cloud.len()) + .map(|idx| { + SortEntry { + key: 1, + index: idx as u32, + } + }) + .collect(); + // TODO: move gaussian_cloud and sorted_entry assets into an asset bundle + #[cfg(feature = "buffer_storage")] + let sorted_entries = sorted_entries_res.add(SortedEntries { + sorted, + }); + + #[cfg(feature = "buffer_texture")] let sorted_entries = sorted_entries_res.add(SortedEntries { - sorted: (0..cloud.gaussians.len()) - .map(|idx| { - SortEntry { - key: 1, - index: idx as u32, - } - }) - .collect(), + texture: images.add(Image::new( + Extent3d { + width: cloud.len_sqrt_ceil() as u32, + height: cloud.len_sqrt_ceil() as u32, + depth_or_array_layers: 1, + }, + TextureDimension::D2, + bytemuck::cast_slice(sorted.as_slice()).to_vec(), + TextureFormat::Rg32Uint, + )), + sorted, }); commands.entity(entity) @@ -179,8 +226,6 @@ pub struct SortEntry { pub index: u32, } -// TODO: add RenderAssetPlugin for SortedEntries & auto-insert to GaussianCloudBundles if their sort mode is not None -// supports pre-sorting or CPU sorting in main world, initializes the sorting_entry_buffer #[derive( Clone, Asset, @@ -193,6 +238,9 @@ pub struct SortEntry { #[uuid = "ac2f08eb-fa13-ccdd-ea11-51571ea332d5"] pub struct SortedEntries { pub sorted: Vec, + + #[cfg(feature = "buffer_texture")] + pub texture: Handle, } impl RenderAsset for SortedEntries { @@ -219,6 +267,9 @@ impl RenderAsset for SortedEntries { Ok(GpuSortedEntry { sorted_entry_buffer, count, + + #[cfg(feature = "buffer_texture")] + texture: sorted_entries.texture, }) } } @@ -230,4 +281,7 @@ impl RenderAsset for SortedEntries { pub struct GpuSortedEntry { pub sorted_entry_buffer: Buffer, pub count: usize, + + #[cfg(feature = "buffer_texture")] + pub texture: Handle, } diff --git a/src/sort/radix.rs b/src/sort/radix.rs index fe75e00f..85003014 100644 --- a/src/sort/radix.rs +++ b/src/sort/radix.rs @@ -9,7 +9,28 @@ use bevy::{ core_pipeline::core_3d::CORE_3D, render::{ render_asset::RenderAssets, - render_resource::*, + render_resource::{ + BindGroup, + BindGroupEntry, + BindGroupLayout, + BindGroupLayoutDescriptor, + BindGroupLayoutEntry, + BindingResource, + BindingType, + Buffer, + BufferBindingType, + BufferDescriptor, + BufferInitDescriptor, + BufferBinding, + BufferSize, + BufferUsages, + CachedComputePipelineId, + CachedPipelineState, + ComputePassDescriptor, + ComputePipelineDescriptor, + PipelineCache, + ShaderStages, + }, renderer::{ RenderContext, RenderDevice, @@ -26,13 +47,15 @@ use bevy::{ view::ViewUniformOffset, }, }; +use static_assertions::assert_cfg; use crate::{ - gaussian::GaussianCloud, + gaussian::cloud::GaussianCloud, GaussianCloudSettings, render::{ GaussianCloudBindGroup, GaussianCloudPipeline, + GaussianCloudPipelineKey, GaussianUniformBindGroups, GaussianViewBindGroup, ShaderDefines, @@ -46,6 +69,15 @@ use crate::{ }; +assert_cfg!( + not(all( + feature = "sort_radix", + feature = "buffer_texture", + )), + "sort_radix and buffer_texture are incompatible", +); + + const RADIX_SHADER_HANDLE: Handle = Handle::weak_from_u128(6234673214); const TEMPORAL_SORT_SHADER_HANDLE: Handle = Handle::weak_from_u128(1634543224); @@ -274,7 +306,7 @@ impl FromWorld for RadixSortPipeline { gaussian_cloud_pipeline.gaussian_cloud_layout.clone(), radix_sort_layout.clone(), ]; - let shader_defs = shader_defs(false, false, false); + let shader_defs = shader_defs(GaussianCloudPipelineKey::default()); let pipeline_cache = render_world.resource::(); let radix_sort_a = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { @@ -322,6 +354,7 @@ pub struct RadixBindGroup { pub radix_sort_bind_groups: [BindGroup; 4], } +#[allow(clippy::too_many_arguments)] pub fn queue_radix_bind_group( mut commands: Commands, radix_pipeline: Res, @@ -425,7 +458,7 @@ pub fn queue_radix_bind_group( &sorting_assets.entry_buffer_b }, offset: 0, - size: BufferSize::new((cloud.count as usize * std::mem::size_of::()) as u64), + size: BufferSize::new((cloud.count * std::mem::size_of::()) as u64), }), }, BindGroupEntry { @@ -437,7 +470,7 @@ pub fn queue_radix_bind_group( &sorted_entries.sorted_entry_buffer }, offset: 0, - size: BufferSize::new((cloud.count as usize * std::mem::size_of::()) as u64), + size: BufferSize::new((cloud.count * std::mem::size_of::()) as u64), }), }, ], @@ -608,7 +641,7 @@ impl Node for RadixSortNode { let mut pass = command_encoder.begin_compute_pass(&ComputePassDescriptor::default()); let radix_sort_c = pipeline_cache.get_compute_pipeline(pipeline.radix_sort_pipelines[2]).unwrap(); - pass.set_pipeline(&radix_sort_c); + pass.set_pipeline(radix_sort_c); pass.set_bind_group( 0, diff --git a/src/sort/radix.wgsl b/src/sort/radix.wgsl index 3b929fc6..99e7c90d 100644 --- a/src/sort/radix.wgsl +++ b/src/sort/radix.wgsl @@ -2,7 +2,6 @@ view, globals, gaussian_uniforms, - points, sorting_pass_index, sorting, status_counters, @@ -18,6 +17,18 @@ in_frustum, } +#ifdef PACKED_F32 +#import bevy_gaussian_splatting::packed::get_position +#endif + +#ifdef BUFFER_STORAGE +#import bevy_gaussian_splatting::planar::get_position +#endif + +#ifdef BUFFER_TEXTURE +#import bevy_gaussian_splatting::texture::get_position +#endif + struct SortingGlobal { digit_histogram: array, #{RADIX_BASE}>, #{RADIX_DIGIT_PLACES}>, @@ -49,11 +60,12 @@ fn radix_sort_a( let start_entry_index = thread_index * #{ENTRIES_PER_INVOCATION_A}u; let end_entry_index = start_entry_index + #{ENTRIES_PER_INVOCATION_A}u; for(var entry_index = start_entry_index; entry_index < end_entry_index; entry_index += 1u) { - if(entry_index >= arrayLength(&points)) { + if(entry_index >= gaussian_uniforms.count) { continue; } var key: u32 = 0xFFFFFFFFu; // Stream compaction for frustum culling - let transformed_position = (gaussian_uniforms.global_transform * points[entry_index].position).xyz; + let position = vec4(get_position(entry_index), 1.0); + let transformed_position = (gaussian_uniforms.global_transform * position).xyz; let clip_space_pos = world_to_clip(transformed_position); if(in_frustum(clip_space_pos.xyz)) { // key = bitcast(1.0 - clip_space_pos.z); @@ -158,7 +170,7 @@ fn radix_sort_c( let assignment = sorting_shared_c.entries[0]; let global_entry_offset = assignment * #{WORKGROUP_ENTRIES_C}u; // TODO: Specialize end shader - if(gl_LocalInvocationID.x == 0u && assignment * #{WORKGROUP_ENTRIES_C}u + #{WORKGROUP_ENTRIES_C}u >= arrayLength(&points)) { + if(gl_LocalInvocationID.x == 0u && assignment * #{WORKGROUP_ENTRIES_C}u + #{WORKGROUP_ENTRIES_C}u >= gaussian_uniforms.count) { // Last workgroup resets the assignment number for the next pass sorting.assignment_counter = 0u; } @@ -199,7 +211,7 @@ fn radix_sort_c( } } atomicStore(&status_counters[assignment][gl_LocalInvocationID.x], 0x80000000u | (global_digit_count + local_digit_count)); - if(sorting_pass_index == #{RADIX_DIGIT_PLACES}u - 1u && gl_LocalInvocationID.x == #{WORKGROUP_INVOCATIONS_C}u - 2u && global_entry_offset + #{WORKGROUP_ENTRIES_C}u >= arrayLength(&points)) { + if(sorting_pass_index == #{RADIX_DIGIT_PLACES}u - 1u && gl_LocalInvocationID.x == #{WORKGROUP_INVOCATIONS_C}u - 2u && global_entry_offset + #{WORKGROUP_ENTRIES_C}u >= gaussian_uniforms.count) { draw_indirect.vertex_count = 4u; draw_indirect.instance_count = global_digit_count + local_digit_count; } diff --git a/src/sort/rayon.rs b/src/sort/rayon.rs index 25f1a301..2edce7f2 100644 --- a/src/sort/rayon.rs +++ b/src/sort/rayon.rs @@ -40,25 +40,41 @@ pub fn rayon_sort( )>, mut last_camera_position: Local, mut last_sort_time: Local>, + mut period: Local, + mut sort_done: Local, ) { - let period = std::time::Duration::from_millis(100); + if last_sort_time.is_none() { + *period = std::time::Duration::from_millis(100); + } + if let Some(last_sort_time) = last_sort_time.as_ref() { - if last_sort_time.elapsed() < period { + if last_sort_time.elapsed() < *period { return; } } // TODO: move sort to render world, use extracted views and update the existing buffer instead of creating new + let sort_start_time = Instant::now(); + let mut performed_sort = false; + for ( camera_transform, _camera, ) in cameras.iter() { let camera_position = camera_transform.compute_transform().translation; - if *last_camera_position == camera_position { - return; + let camera_movement = *last_camera_position != camera_position; + + if camera_movement { + *sort_done = false; + } else { + if *sort_done { + return; + } } + *last_camera_position = camera_position; + for ( gaussian_cloud_handle, sorted_entries_handle, @@ -78,16 +94,18 @@ pub fn rayon_sort( if let Some(gaussian_cloud) = gaussian_clouds_res.get(gaussian_cloud_handle) { if let Some(sorted_entries) = sorted_entries_res.get_mut(sorted_entries_handle) { - assert_eq!(gaussian_cloud.gaussians.len(), sorted_entries.sorted.len()); + assert_eq!(gaussian_cloud.len(), sorted_entries.sorted.len()); - *last_camera_position = camera_position; + *sort_done = true; *last_sort_time = Some(Instant::now()); - gaussian_cloud.gaussians.par_iter() + performed_sort = true; + + gaussian_cloud.position_par_iter() .zip(sorted_entries.sorted.par_iter_mut()) .enumerate() - .for_each(|(idx, (gaussian, sort_entry))| { - let position = Vec3::from_slice(gaussian.position.as_ref()); + .for_each(|(idx, (position, sort_entry))| { + let position = Vec3::from_slice(position.as_ref()); let delta = camera_position - position; sort_entry.key = bytemuck::cast(delta.length_squared()); @@ -103,4 +121,15 @@ pub fn rayon_sort( } } } + + let sort_end_time = Instant::now(); + let delta = sort_end_time - sort_start_time; + + if performed_sort { + *period = std::time::Duration::from_millis( + 100 + .max(period.as_millis() as u64 * 4 / 5) + .max(10 * delta.as_millis() as u64) + ); + } } diff --git a/src/sort/std.rs b/src/sort/std.rs index 3bafa2fa..1f9dc6aa 100644 --- a/src/sort/std.rs +++ b/src/sort/std.rs @@ -23,6 +23,7 @@ impl Plugin for StdSortPlugin { } } +// TODO: async CPU sort to prevent frame drops on large clouds pub fn std_sort( asset_server: Res, gaussian_clouds_res: Res>, @@ -38,20 +39,42 @@ pub fn std_sort( )>, mut last_camera_position: Local, mut last_sort_time: Local>, + mut period: Local, + mut camera_debounce: Local, + mut sort_done: Local, ) { - let period = std::time::Duration::from_millis(100); + if last_sort_time.is_none() { + *period = std::time::Duration::from_millis(100); + } + if let Some(last_sort_time) = last_sort_time.as_ref() { - if last_sort_time.elapsed() < period { + if last_sort_time.elapsed() < *period { return; } } + let sort_start_time = Instant::now(); + let mut performed_sort = false; + for ( camera_transform, _camera, ) in cameras.iter() { let camera_position = camera_transform.compute_transform().translation; - if *last_camera_position == camera_position { + let camera_movement = *last_camera_position != camera_position; + + if camera_movement { + *sort_done = false; + *camera_debounce = true; + } else { + if *sort_done { + return; + } + } + + if *camera_debounce { + *last_camera_position = camera_position; + *camera_debounce = false; return; } @@ -74,16 +97,18 @@ pub fn std_sort( if let Some(gaussian_cloud) = gaussian_clouds_res.get(gaussian_cloud_handle) { if let Some(sorted_entries) = sorted_entries_res.get_mut(sorted_entries_handle) { - assert_eq!(gaussian_cloud.gaussians.len(), sorted_entries.sorted.len()); + assert_eq!(gaussian_cloud.len(), sorted_entries.sorted.len()); - *last_camera_position = camera_position; + *sort_done = true; *last_sort_time = Some(Instant::now()); - gaussian_cloud.gaussians.iter() + performed_sort = true; + + gaussian_cloud.position_iter() .zip(sorted_entries.sorted.iter_mut()) .enumerate() - .for_each(|(idx, (gaussian, sort_entry))| { - let position = Vec3::from_slice(gaussian.position.as_ref()); + .for_each(|(idx, (position, sort_entry))| { + let position = Vec3::from_slice(position.as_ref()); let delta = camera_position - position; sort_entry.key = bytemuck::cast(delta.length_squared()); @@ -99,4 +124,15 @@ pub fn std_sort( } } } + + let sort_end_time = Instant::now(); + let delta = sort_end_time - sort_start_time; + + if performed_sort { + *period = std::time::Duration::from_millis( + 100 + .max(period.as_millis() as u64 * 4 / 5) + .max(10 * delta.as_millis() as u64) + ); + } } diff --git a/src/sort/temporal.wgsl b/src/sort/temporal.wgsl index 894dcc6d..17ebd8be 100644 --- a/src/sort/temporal.wgsl +++ b/src/sort/temporal.wgsl @@ -19,7 +19,7 @@ fn temporal_sort_flop( // // pair sort entries in window size // for (var i = start_index; i < end_index; i += 2u) { - // let pos_a = points[input_entries[i][0]].position.xyz; + // let pos_a = points[input_entries[i][0]].position_visibility.xyz; // let depth_a = world_to_clip(pos_a).z; // } } diff --git a/src/utils.rs b/src/utils.rs index 24b72910..f03ff538 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -23,9 +23,9 @@ pub fn get_arg(n: usize) -> Option { #[cfg(target_arch = "wasm32")] pub fn get_arg(n: usize) -> Option { - let window = web_sys::window()?; // Get the window object - let location = window.location(); // Get the location object from the window - let search = location.search().ok()?; // Get the search (query string) part of the URL + let window = web_sys::window()?; + let location = window.location(); + let search = location.search().ok()?; let args = search .trim_start_matches('?') diff --git a/tests/gpu/_harness.rs b/tests/gpu/_harness.rs index 81821b43..f96e9834 100644 --- a/tests/gpu/_harness.rs +++ b/tests/gpu/_harness.rs @@ -45,18 +45,10 @@ pub fn test_harness_app( } +#[derive(Default)] pub struct TestState { pub test_loaded: bool, pub test_completed: bool, } -impl Default for TestState { - fn default() -> Self { - TestState { - test_loaded: false, - test_completed: false, - } - } -} - pub type TestStateArc = Arc>; diff --git a/tests/gpu/gaussian.rs b/tests/gpu/gaussian.rs index 03f0ba99..bdd19d4a 100644 --- a/tests/gpu/gaussian.rs +++ b/tests/gpu/gaussian.rs @@ -77,20 +77,16 @@ fn check_image_equality(image: &Image, other: &Image) -> bool { fn test_stability(captures: Arc>>) { let all_frames_similar = captures.lock().unwrap().iter() - .fold(Some(None), |acc, image| { + .try_fold(None, |acc, image| { match acc { Some(acc_image) => { - if let Some(acc_image) = acc_image { - if check_image_equality(acc_image, image) { - Some(Some(acc_image)) - } else { - None - } + if check_image_equality(acc_image, image) { + Some(Some(acc_image)) } else { - Some(Some(image)) + None } }, - None => None, + None => Some(Some(image)), } }).is_some(); assert!(all_frames_similar, "all frames are not the same"); @@ -153,7 +149,7 @@ fn capture_ready( if let Ok(window_entity) = main_window.get_single() { screenshot_manager.take_screenshot(window_entity, move |image: Image| { - let has_non_zero_data = image.data.iter().fold(false, |non_zero, &x| non_zero || x != 0); + let has_non_zero_data = image.data.iter().any(|&x| x != 0); assert!(has_non_zero_data, "screenshot is all zeros"); let mut buffer = buffer_clone.lock().unwrap(); diff --git a/tests/gpu/radix.rs b/tests/gpu/radix.rs index d448099f..2ad28548 100644 --- a/tests/gpu/radix.rs +++ b/tests/gpu/radix.rs @@ -206,7 +206,7 @@ impl Node for RadixTestNode { return depth_acc; } - let position = gaussians[idx].position; + let position = gaussians[idx].position_visibility; let position_vec3 = Vec3::new(position[0], position[1], position[2]); let depth = (position_vec3 - camera_position).length(); diff --git a/tests/io.rs b/tests/io.rs index 5cb2a9fd..d6469b62 100644 --- a/tests/io.rs +++ b/tests/io.rs @@ -7,7 +7,7 @@ use bevy_gaussian_splatting::{ #[test] fn test_codec() { - let count = 100; + let count = 10000; let gaussians = random_gaussians(count); let encoded = gaussians.encode(); diff --git a/tools/compare_aabb_obb.rs b/tools/compare_aabb_obb.rs index d48dea99..279cec82 100644 --- a/tools/compare_aabb_obb.rs +++ b/tools/compare_aabb_obb.rs @@ -49,24 +49,22 @@ pub fn setup_aabb_obb_compare( mut gaussian_assets: ResMut>, ) { let mut blue_sh = SphericalHarmonicCoefficients::default(); - blue_sh.coefficients[2] = 5.0; + blue_sh.set(2, 5.0); let blue_aabb_gaussian = Gaussian { - position: [0.0, 0.0, 0.0, 1.0], - rotation: [0.89, 0.0, -0.432, 0.144], - scale_opacity: [10.0, 1.0, 1.0, 0.5], + position_visibility: [0.0, 0.0, 0.0, 1.0].into(), + rotation: [0.89, 0.0, -0.432, 0.144].into(), + scale_opacity: [10.0, 1.0, 1.0, 0.5].into(), spherical_harmonic: blue_sh, }; commands.spawn(( GaussianSplattingBundle { cloud: gaussian_assets.add( - GaussianCloud { - gaussians: vec![ - blue_aabb_gaussian, - blue_aabb_gaussian, - ], - } + GaussianCloud::from_gaussians(vec![ + blue_aabb_gaussian, + blue_aabb_gaussian, + ]) ), settings: GaussianCloudSettings { aabb: true, @@ -79,24 +77,22 @@ pub fn setup_aabb_obb_compare( )); let mut red_sh = SphericalHarmonicCoefficients::default(); - red_sh.coefficients[0] = 5.0; + red_sh.set(0, 5.0); let red_obb_gaussian = Gaussian { - position: [0.0, 0.0, 0.0, 1.0], - rotation: [0.89, 0.0, -0.432, 0.144], - scale_opacity: [10.0, 1.0, 1.0, 0.5], + position_visibility: [0.0, 0.0, 0.0, 1.0].into(), + rotation: [0.89, 0.0, -0.432, 0.144].into(), + scale_opacity: [10.0, 1.0, 1.0, 0.5].into(), spherical_harmonic: red_sh, }; commands.spawn(( GaussianSplattingBundle { cloud: gaussian_assets.add( - GaussianCloud { - gaussians: vec![ - red_obb_gaussian, - red_obb_gaussian, - ], - } + GaussianCloud::from_gaussians(vec![ + red_obb_gaussian, + red_obb_gaussian, + ]) ), settings: GaussianCloudSettings { aabb: false, diff --git a/tools/ply_to_gcloud.rs b/tools/ply_to_gcloud.rs index 3f3cbedd..4cad1a82 100644 --- a/tools/ply_to_gcloud.rs +++ b/tools/ply_to_gcloud.rs @@ -1,5 +1,3 @@ -use std::io::Write; - use byte_unit::{ Byte, UnitType, @@ -8,11 +6,32 @@ use byte_unit::{ use bevy_gaussian_splatting::{ GaussianCloud, io::{ - codec::GaussianCloudCodec, ply::parse_ply, + writer::write_gaussian_cloud_to_file, }, }; +#[cfg(feature = "query_sparse")] +use bevy_gaussian_splatting::query::sparse::SparseSelect; + + +#[allow(dead_code)] +fn is_point_in_transformed_sphere(pos: &[f32; 3]) -> bool { + let inv_scale_x = 1.0 / 1.75; + let inv_scale_y = 1.0 / 1.75; + let inv_scale_z = 1.0 / 1.75; + + let inv_trans_x = 1.7; + let inv_trans_y = -0.5; + let inv_trans_z = -3.8; + + let transformed_x = (pos[0] + inv_trans_x) * inv_scale_x; + let transformed_y = (pos[1] + inv_trans_y) * inv_scale_y; + let transformed_z = (pos[2] + inv_trans_z) * inv_scale_z; + + transformed_x.powi(2) + transformed_y.powi(2) + transformed_z.powi(2) <= 1.0 +} + fn main() { let filename = std::env::args().nth(1).expect("no filename given"); @@ -22,20 +41,37 @@ fn main() { let file = std::fs::File::open(&filename).expect("failed to open file"); let mut reader = std::io::BufReader::new(file); - let cloud = GaussianCloud { - gaussians: parse_ply(&mut reader).expect("failed to parse ply file"), - ..Default::default() - }; + let mut cloud = GaussianCloud::from_gaussians( + parse_ply(&mut reader).expect("failed to parse ply file"), + ); + + // TODO: prioritize mesh selection over export filter + // println!("initial cloud size: {}", cloud.len()); + // cloud = (0..cloud.len()) + // .filter(|&idx| { + // is_point_in_transformed_sphere( + // cloud.position(idx), + // ) + // }) + // .map(|idx| cloud.gaussian(idx)) + // .collect(); + // println!("filtered position cloud size: {}", cloud.len()); + + #[cfg(feature = "query_sparse")] + { + let sparse_selection = SparseSelect::default().select(&cloud).invert(cloud.len()); + + cloud = sparse_selection.indicies.iter() + .map(|idx| cloud.gaussian(*idx)) + .collect(); + println!("sparsity filtered cloud size: {}", cloud.len()); + } let base_filename = filename.split('.').next().expect("no extension").to_string(); let gcloud_filename = base_filename + ".gcloud"; - let gcloud_file = std::fs::File::create(&gcloud_filename).expect("failed to create file"); - let mut gcloud_writer = std::io::BufWriter::new(gcloud_file); - - let data = cloud.encode(); - gcloud_writer.write_all(data.as_slice()).expect("failed to write to gcloud file"); + write_gaussian_cloud_to_file(&cloud, &gcloud_filename); let post_encode_bytes = Byte::from_u64(std::fs::metadata(&gcloud_filename).expect("failed to get metadata").len()); - println!("output file size: {}", post_encode_bytes.get_appropriate_unit(UnitType::Decimal).to_string()); + println!("output file size: {}", post_encode_bytes.get_appropriate_unit(UnitType::Decimal)); } diff --git a/viewer/viewer.rs b/viewer/viewer.rs index aa03ad84..4ebb1e2b 100644 --- a/viewer/viewer.rs +++ b/viewer/viewer.rs @@ -25,12 +25,24 @@ use bevy_gaussian_splatting::{ }, }; +#[cfg(feature = "material_noise")] +use bevy_gaussian_splatting::material::noise::NoiseMaterial; + #[cfg(feature = "morph_particles")] use bevy_gaussian_splatting::morph::particle::{ ParticleBehaviors, random_particle_behaviors, }; +#[cfg(feature = "query_select")] +use bevy_gaussian_splatting::query::select::{ + InvertSelectionEvent, + SaveSelectionEvent, +}; + +#[cfg(feature = "query_sparse")] +use bevy_gaussian_splatting::query::sparse::SparseSelect; + pub struct GaussianSplattingViewer { pub editor: bool, @@ -68,7 +80,7 @@ fn setup_gaussian_cloud( println!("generating {} gaussians", n); cloud = gaussian_assets.add(random_gaussians(n)); } else if let Some(filename) = file_arg { - if filename == "--help".to_string() { + if filename == "--help" { println!("usage: cargo run -- [filename | n]"); return; } @@ -95,6 +107,9 @@ fn setup_gaussian_cloud( }, PanOrbitCamera{ allow_upside_down: true, + orbit_smoothness: 0.0, + pan_smoothness: 0.0, + zoom_smoothness: 0.0, ..default() }, )); @@ -132,6 +147,78 @@ fn setup_particle_behavior( } } +#[cfg(feature = "material_noise")] +fn setup_noise_material( + mut commands: Commands, + asset_server: Res, + gaussian_clouds: Query<( + Entity, + &Handle, + Without, + )>, +) { + if gaussian_clouds.is_empty() { + return; + } + + for ( + entity, + cloud_handle, + _ + ) in gaussian_clouds.iter() { + if Some(bevy::asset::LoadState::Loading) == asset_server.get_load_state(cloud_handle) { + continue; + } + + commands.entity(entity) + .insert(NoiseMaterial::default()); + } +} + + +#[cfg(feature = "query_select")] +fn press_i_invert_selection( + keys: Res>, + mut select_inverse_events: EventWriter, +) { + if keys.just_pressed(KeyCode::I) { + println!("inverting selection"); + select_inverse_events.send(InvertSelectionEvent); + } +} + +#[cfg(feature = "query_select")] +fn press_o_save_selection( + keys: Res>, + mut select_inverse_events: EventWriter, +) { + if keys.just_pressed(KeyCode::O) { + println!("saving selection"); + select_inverse_events.send(SaveSelectionEvent); + } +} + + +#[cfg(feature = "query_sparse")] +fn setup_sparse_select( + mut commands: Commands, + gaussian_cloud: Query<( + Entity, + &Handle, + Without, + )>, +) { + if gaussian_cloud.is_empty() { + return; + } + + commands.entity(gaussian_cloud.single().0) + .insert(SparseSelect { + completed: true, + ..default() + }); +} + fn example_app() { let config = GaussianSplattingViewer::default(); @@ -181,7 +268,7 @@ fn example_app() { } if config.show_fps { - app.add_plugins(FrameTimeDiagnosticsPlugin::default()); + app.add_plugins(FrameTimeDiagnosticsPlugin); app.add_systems(Startup, fps_display_setup); app.add_systems(Update, fps_update_system); } @@ -191,9 +278,21 @@ fn example_app() { app.add_plugins(GaussianSplattingPlugin); app.add_systems(Startup, setup_gaussian_cloud); + #[cfg(feature = "material_noise")] + app.add_systems(Update, setup_noise_material); + #[cfg(feature = "morph_particles")] app.add_systems(Update, setup_particle_behavior); + #[cfg(feature = "query_select")] + { + app.add_systems(Update, press_i_invert_selection); + app.add_systems(Update, press_o_save_selection); + } + + #[cfg(feature = "query_sparse")] + app.add_systems(Update, setup_sparse_select); + app.run(); } diff --git a/www/README.md b/www/README.md index 5e5d7d68..114e8701 100644 --- a/www/README.md +++ b/www/README.md @@ -5,7 +5,7 @@ to build wasm run: ```bash -cargo build --target wasm32-unknown-unknown --release --no-default-features --features "io_flexbuffers sort_std viewer" +cargo build --target wasm32-unknown-unknown --release --no-default-features --features "web" ``` to generate bindings: