From 8d753872e9a62450856429bd8731c8bc9b0127cf Mon Sep 17 00:00:00 2001 From: Mitchell Mosure Date: Fri, 15 Dec 2023 19:35:50 -0600 Subject: [PATCH] feat: depth visualization (#55) --- Cargo.toml | 2 +- README.md | 3 ++- src/gaussian.rs | 2 ++ src/morph/particle.rs | 2 +- src/render/bindings.wgsl | 1 + src/render/color.wgsl | 12 ++++++++++++ src/render/gaussian.wgsl | 25 ++++++++++++++++++++++++- src/render/mod.rs | 36 +++++++++++++++++++++++++++++++++--- src/render/transform.wgsl | 1 - src/sort/radix.rs | 2 +- 10 files changed, 77 insertions(+), 9 deletions(-) create mode 100644 src/render/color.wgsl diff --git a/Cargo.toml b/Cargo.toml index eaf9d398..11772fa7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "bevy_gaussian_splatting" description = "bevy gaussian splatting render pipeline plugin" -version = "0.5.0" +version = "1.0.0" edition = "2021" authors = ["mosure "] license = "MIT" diff --git a/README.md b/README.md index e203f62c..ab2a8507 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ download [cactus.gcloud](https://mitchell.mosure.me/cactus.gcloud) - [X] bevy gaussian cloud render pipeline - [X] gaussian cloud particle effects - [X] wasm support /w [live demo](https://mosure.github.io/bevy_gaussian_splatting/index.html?arg1=icecream.gcloud) +- [X] depth colorization - [ ] 4D gaussian cloud wavelet compression - [ ] accelerated spatial queries - [ ] temporal depth sorting @@ -74,7 +75,7 @@ fn setup_gaussian_cloud( | `bevy_gaussian_splatting` | `bevy` | | :-- | :-- | -| `0.4 - 0.5` | `0.12` | +| `0.4 - 1.0` | `0.12` | | `0.1 - 0.3` | `0.11` | diff --git a/src/gaussian.rs b/src/gaussian.rs index c83b6c72..a7bd33a1 100644 --- a/src/gaussian.rs +++ b/src/gaussian.rs @@ -216,6 +216,7 @@ pub struct GaussianCloudSettings { pub global_scale: f32, pub global_transform: GlobalTransform, pub visualize_bounding_box: bool, + pub visualize_depth: bool, pub sort_mode: SortMode, } @@ -226,6 +227,7 @@ impl Default for GaussianCloudSettings { global_scale: 2.0, global_transform: Transform::IDENTITY.into(), visualize_bounding_box: false, + visualize_depth: false, sort_mode: SortMode::default(), } } diff --git a/src/morph/particle.rs b/src/morph/particle.rs index e728e042..9a68d875 100644 --- a/src/morph/particle.rs +++ b/src/morph/particle.rs @@ -201,7 +201,7 @@ impl FromWorld for ParticleBehaviorPipeline { ], }); - let shader_defs = shader_defs(false, false); + let shader_defs = shader_defs(false, false, false); let pipeline_cache = render_world.resource::(); let particle_behavior_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { diff --git a/src/render/bindings.wgsl b/src/render/bindings.wgsl index 8e641b09..9f9f47c5 100644 --- a/src/render/bindings.wgsl +++ b/src/render/bindings.wgsl @@ -10,6 +10,7 @@ struct GaussianUniforms { global_transform: mat4x4, global_scale: f32, + count: u32, }; @group(1) @binding(0) var gaussian_uniforms: GaussianUniforms; diff --git a/src/render/color.wgsl b/src/render/color.wgsl new file mode 100644 index 00000000..d640dc55 --- /dev/null +++ b/src/render/color.wgsl @@ -0,0 +1,12 @@ +#define_import_path bevy_gaussian_splatting::color + + +fn depth_to_rgb(depth: f32, min_depth: f32, max_depth: f32) -> vec3 { + let normalized_depth = clamp((depth - min_depth) / (max_depth - min_depth), 0.0, 1.0); + + let r = smoothstep(0.5, 1.0, normalized_depth); + let g = 1.0 - abs(normalized_depth - 0.5) * 2.0; + let b = 1.0 - smoothstep(0.0, 0.5, normalized_depth); + + return vec3(r, g, b); +} diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index 1c5c7228..5061c968 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -10,6 +10,9 @@ output_entries, Entry, } +#import bevy_gaussian_splatting::color::{ + depth_to_rgb, +} #import bevy_gaussian_splatting::spherical_harmonics::spherical_harmonics_lookup #import bevy_gaussian_splatting::transform::{ world_to_clip, @@ -233,8 +236,28 @@ fn vs_points( let quad_offset = quad_vertices[quad_index]; let ray_direction = normalize(transformed_position - view.world_position); + +#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 camera_position = view.world_position; + + let min_distance = length(min_position - camera_position); + let max_distance = length(max_position - camera_position); + + let depth = length(transformed_position - camera_position); + let rgb = depth_to_rgb( + depth, + min_distance, + max_distance, + ); +#else + let rgb = spherical_harmonics_lookup(ray_direction, point.sh); +#endif + output.color = vec4( - spherical_harmonics_lookup(ray_direction, point.sh), + rgb, point.scale_opacity.a ); diff --git a/src/render/mod.rs b/src/render/mod.rs index ce2d0238..4fb60409 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -71,6 +71,7 @@ use crate::{ 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 TRANSFORM_SHADER_HANDLE: Handle = Handle::weak_from_u128(734523534); @@ -88,6 +89,13 @@ impl Plugin for RenderPipelinePlugin { Shader::from_wgsl ); + load_internal_asset!( + app, + COLOR_SHADER_HANDLE, + "color.wgsl", + Shader::from_wgsl + ); + load_internal_asset!( app, GAUSSIAN_SHADER_HANDLE, @@ -147,7 +155,7 @@ pub struct GpuGaussianSplattingBundle { pub settings: GaussianCloudSettings, pub settings_uniform: GaussianCloudUniform, pub sorted_entries: Handle, - pub verticies: Handle, + pub cloud_handle: Handle, } #[derive(Debug, Clone)] @@ -253,6 +261,7 @@ fn queue_gaussians( let key = GaussianCloudPipelineKey { aabb: settings.aabb, visualize_bounding_box: settings.visualize_bounding_box, + visualize_depth: settings.visualize_depth, }; let pipeline = pipelines.specialize(&pipeline_cache, &custom_pipeline, key); @@ -443,6 +452,7 @@ impl Default for ShaderDefines { pub fn shader_defs( aabb: bool, visualize_bounding_box: bool, + visualize_depth: bool, ) -> Vec { let defines = ShaderDefines::default(); let mut shader_defs = vec![ @@ -474,6 +484,10 @@ pub fn shader_defs( #[cfg(feature = "morph_particles")] shader_defs.push("READ_WRITE_POINTS".into()); + if visualize_depth { + shader_defs.push("VISUALIZE_DEPTH".into()); + } + shader_defs } @@ -481,6 +495,7 @@ pub fn shader_defs( pub struct GaussianCloudPipelineKey { pub aabb: bool, pub visualize_bounding_box: bool, + pub visualize_depth: bool, } impl SpecializedRenderPipeline for GaussianCloudPipeline { @@ -490,6 +505,7 @@ impl SpecializedRenderPipeline for GaussianCloudPipeline { let shader_defs = shader_defs( key.aabb, key.visualize_bounding_box, + key.visualize_depth, ); RenderPipelineDescriptor { @@ -563,11 +579,14 @@ type DrawGaussians = ( pub struct GaussianCloudUniform { pub transform: Mat4, pub global_scale: f32, + pub count: u32, } pub fn extract_gaussians( mut commands: Commands, mut prev_commands_len: Local, + asset_server: Res, + gaussian_cloud_res: Res>, gaussians_query: Extract< Query<( Entity, @@ -585,7 +604,7 @@ pub fn extract_gaussians( for ( entity, visibility, - verticies, + cloud_handle, sorted_entries, settings, ) in gaussians_query.iter() { @@ -593,9 +612,20 @@ pub fn extract_gaussians( continue; } + 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 settings_uniform = GaussianCloudUniform { transform: settings.global_transform.compute_matrix(), global_scale: settings.global_scale, + count: cloud.count as u32, }; commands_list.push(( entity, @@ -603,7 +633,7 @@ pub fn extract_gaussians( settings: settings.clone(), settings_uniform, sorted_entries: sorted_entries.clone(), - verticies: verticies.clone(), + cloud_handle: cloud_handle.clone(), }, )); } diff --git a/src/render/transform.wgsl b/src/render/transform.wgsl index 50d13ebf..5f9b943b 100644 --- a/src/render/transform.wgsl +++ b/src/render/transform.wgsl @@ -8,7 +8,6 @@ fn world_to_clip(world_pos: vec3) -> vec4 { return homogenous_pos / (homogenous_pos.w + 0.000000001); } - // fn world_to_clip(world_pos: vec3) -> vec4 { // let homogenous_pos = view.unjittered_view_proj * vec4(world_pos, 1.0); // return homogenous_pos / (homogenous_pos.w + 0.000000001); diff --git a/src/sort/radix.rs b/src/sort/radix.rs index b9a1f92f..fe75e00f 100644 --- a/src/sort/radix.rs +++ b/src/sort/radix.rs @@ -274,7 +274,7 @@ impl FromWorld for RadixSortPipeline { gaussian_cloud_pipeline.gaussian_cloud_layout.clone(), radix_sort_layout.clone(), ]; - let shader_defs = shader_defs(false, false); + let shader_defs = shader_defs(false, false, false); let pipeline_cache = render_world.resource::(); let radix_sort_a = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {