Skip to content

Commit

Permalink
feat: depth visualization (#55)
Browse files Browse the repository at this point in the history
  • Loading branch information
mosure authored Dec 16, 2023
1 parent 6033195 commit 8d75387
Show file tree
Hide file tree
Showing 10 changed files with 77 additions and 9 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
3 changes: 2 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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` |


Expand Down
2 changes: 2 additions & 0 deletions src/gaussian.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}

Expand All @@ -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(),
}
}
Expand Down
2 changes: 1 addition & 1 deletion src/morph/particle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
1 change: 1 addition & 0 deletions src/render/bindings.wgsl
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
struct GaussianUniforms {
global_transform: mat4x4<f32>,
global_scale: f32,
count: u32,
};
@group(1) @binding(0) var<uniform> gaussian_uniforms: GaussianUniforms;

Expand Down
12 changes: 12 additions & 0 deletions src/render/color.wgsl
Original file line number Diff line number Diff line change
@@ -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);
}
25 changes: 24 additions & 1 deletion src/render/gaussian.wgsl
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
);

Expand Down
36 changes: 33 additions & 3 deletions src/render/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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,
Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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![
Expand Down Expand Up @@ -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 {
Expand All @@ -490,6 +505,7 @@ impl SpecializedRenderPipeline for GaussianCloudPipeline {
let shader_defs = shader_defs(
key.aabb,
key.visualize_bounding_box,
key.visualize_depth,
);

RenderPipelineDescriptor {
Expand Down Expand Up @@ -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,
Expand All @@ -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(),
},
));
}
Expand Down
1 change: 0 additions & 1 deletion src/render/transform.wgsl
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion src/sort/radix.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit 8d75387

Please sign in to comment.