feat: depth visualization (#55)

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 <[email protected]>"]
 license = "MIT"

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` |
 
 

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(),
         }
     }

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::<PipelineCache>();
 
         let particle_behavior_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {

src/render/bindings.wgsl

@@ -10,6 +10,7 @@
 struct GaussianUniforms {
     global_transform: mat4x4<f32>,
     global_scale: f32,
+    count: u32,
 };
 @group(1) @binding(0) var<uniform> gaussian_uniforms: GaussianUniforms;
 

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<f32> {
+    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<f32>(r, g, 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<f32>(
-        spherical_harmonics_lookup(ray_direction, point.sh),
+        rgb,
         point.scale_opacity.a
     );
 

src/render/mod.rs

@@ -71,6 +71,7 @@ use crate::{
 
 
 const BINDINGS_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(675257236);
+const COLOR_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(51234253);
 const GAUSSIAN_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(68294581);
 const SPHERICAL_HARMONICS_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(834667312);
 const TRANSFORM_SHADER_HANDLE: Handle<Shader> = 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<SortedEntries>,
-    pub verticies: Handle<GaussianCloud>,
+    pub cloud_handle: Handle<GaussianCloud>,
 }
 
 #[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<ShaderDefVal> {
     let defines = ShaderDefines::default();
     let mut shader_defs = vec![
@@ -474,13 +484,18 @@ 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
 }
 
 #[derive(PartialEq, Eq, Hash, Clone, Copy)]
 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<usize>,
+    asset_server: Res<AssetServer>,
+    gaussian_cloud_res: Res<RenderAssets<GaussianCloud>>,
     gaussians_query: Extract<
         Query<(
             Entity,
@@ -585,25 +604,36 @@ pub fn extract_gaussians(
     for (
         entity,
         visibility,
-        verticies,
+        cloud_handle,
         sorted_entries,
         settings,
     ) in gaussians_query.iter() {
         if visibility == Visibility::Hidden {
             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,
             GpuGaussianSplattingBundle {
                 settings: settings.clone(),
                 settings_uniform,
                 sorted_entries: sorted_entries.clone(),
-                verticies: verticies.clone(),
+                cloud_handle: cloud_handle.clone(),
             },
         ));
     }

src/render/transform.wgsl

@@ -8,7 +8,6 @@ fn world_to_clip(world_pos: vec3<f32>) -> vec4<f32> {
     return homogenous_pos / (homogenous_pos.w + 0.000000001);
 }
 
-
 // fn world_to_clip(world_pos: vec3<f32>) -> vec4<f32> {
 //     let homogenous_pos = view.unjittered_view_proj * vec4<f32>(world_pos, 1.0);
 //     return homogenous_pos / (homogenous_pos.w + 0.000000001);

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::<PipelineCache>();
         let radix_sort_a = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {