From 063d9f0b4256462efcc8902baf36c5a73504ca3c Mon Sep 17 00:00:00 2001 From: Mitchell Mosure Date: Sun, 9 Jun 2024 21:12:50 -0500 Subject: [PATCH] feat: rasterize normal option (#110) --- src/gaussian/settings.rs | 23 +++++++++++-- src/render/gaussian.wgsl | 74 +++++++++++++++++++++++++++++++--------- src/render/mod.rs | 15 ++++---- 3 files changed, 87 insertions(+), 25 deletions(-) diff --git a/src/gaussian/settings.rs b/src/gaussian/settings.rs index 10f6ca5c..42484c2f 100644 --- a/src/gaussian/settings.rs +++ b/src/gaussian/settings.rs @@ -20,6 +20,25 @@ pub enum GaussianCloudDrawMode { HighlightSelected, } + +#[derive( + Clone, + Copy, + Debug, + Default, + Eq, + Hash, + PartialEq, + Reflect, +)] +pub enum GaussianCloudRasterize { + #[default] + Color, + Depth, + Normal, +} + + #[derive(Component, Reflect, Clone)] #[reflect(Component)] pub struct GaussianCloudSettings { @@ -27,9 +46,9 @@ pub struct GaussianCloudSettings { pub global_scale: f32, pub transform: Transform, pub visualize_bounding_box: bool, - pub visualize_depth: bool, pub sort_mode: SortMode, pub draw_mode: GaussianCloudDrawMode, + pub rasterize_mode: GaussianCloudRasterize, } impl Default for GaussianCloudSettings { @@ -39,9 +58,9 @@ impl Default for GaussianCloudSettings { global_scale: 1.0, transform: Transform::IDENTITY, visualize_bounding_box: false, - visualize_depth: false, sort_mode: SortMode::default(), draw_mode: GaussianCloudDrawMode::default(), + rasterize_mode: GaussianCloudRasterize::default(), } } } diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index 77709f37..1e5f8d4e 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -130,27 +130,15 @@ struct GaussianVertexOutput { #endif -// https://github.com/cvlab-epfl/gaussian-splatting-web/blob/905b3c0fb8961e42c79ef97e64609e82383ca1c2/src/shaders.ts#L185 -// TODO: precompute -fn compute_cov3d(scale: vec3, rotation: vec4) -> array { - let S = mat3x3( - scale.x * gaussian_uniforms.global_scale, 0.0, 0.0, - 0.0, scale.y * gaussian_uniforms.global_scale, 0.0, - 0.0, 0.0, scale.z * gaussian_uniforms.global_scale, - ); - +fn get_rotation_matrix( + rotation: vec4, +) -> mat3x3 { let r = rotation.x; let x = rotation.y; let y = rotation.z; let z = rotation.w; - let T = mat3x3( - gaussian_uniforms.transform[0].xyz, - gaussian_uniforms.transform[1].xyz, - gaussian_uniforms.transform[2].xyz, - ); - - let R = mat3x3( + return mat3x3( 1.0 - 2.0 * (y * y + z * z), 2.0 * (x * y - r * z), 2.0 * (x * z + r * y), @@ -163,6 +151,31 @@ fn compute_cov3d(scale: vec3, rotation: vec4) -> array { 2.0 * (y * z + r * x), 1.0 - 2.0 * (x * x + y * y), ); +} + +fn get_scale_matrix( + scale: vec3, +) -> mat3x3 { + return mat3x3( + scale.x * gaussian_uniforms.global_scale, 0.0, 0.0, + 0.0, scale.y * gaussian_uniforms.global_scale, 0.0, + 0.0, 0.0, scale.z * gaussian_uniforms.global_scale, + ); +} + + +// https://github.com/cvlab-epfl/gaussian-splatting-web/blob/905b3c0fb8961e42c79ef97e64609e82383ca1c2/src/shaders.ts#L185 +// TODO: precompute +fn compute_cov3d(scale: vec3, rotation: vec4) -> array { + let S = get_scale_matrix(scale); + + let T = mat3x3( + gaussian_uniforms.transform[0].xyz, + gaussian_uniforms.transform[1].xyz, + gaussian_uniforms.transform[2].xyz, + ); + + let R = get_rotation_matrix(rotation); let M = S * R; let Sigma = transpose(M) * M; @@ -356,7 +369,7 @@ fn vs_points( var rgb = vec3(0.0); -#ifdef VISUALIZE_DEPTH +#ifdef RASTERIZE_DEPTH 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); @@ -374,6 +387,33 @@ fn vs_points( min_distance, max_distance, ); +#else ifdef RASTERIZE_NORMAL + let T = mat3x3( + gaussian_uniforms.transform[0].xyz, + gaussian_uniforms.transform[1].xyz, + gaussian_uniforms.transform[2].xyz, + ); + + let R = get_rotation_matrix(get_rotation(splat_index)); + let S = get_scale_matrix(get_scale(splat_index)); + + let M = S * R; + let Sigma = transpose(M) * M; + + let N = T * Sigma * transpose(T); + let normal = vec3( + N[0][0], + N[0][1], + N[1][1], + ); + + let t = normalize(normal); + + rgb = vec3( + 0.5 * (t.x + 1.0), + 0.5 * (t.y + 1.0), + 0.5 * (t.z + 1.0) + ); #else rgb = get_color(splat_index, ray_direction); #endif diff --git a/src/render/mod.rs b/src/render/mod.rs index ab18358b..6c6bd0d1 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -61,6 +61,7 @@ use crate::{ cloud::GaussianCloud, settings::{ GaussianCloudDrawMode, + GaussianCloudRasterize, GaussianCloudSettings, }, }, @@ -305,8 +306,8 @@ fn queue_gaussians( let key = GaussianCloudPipelineKey { aabb: settings.aabb, visualize_bounding_box: settings.visualize_bounding_box, - visualize_depth: settings.visualize_depth, draw_mode: settings.draw_mode, + rasterize_mode: settings.rasterize_mode, sample_count: msaa.samples(), }; @@ -521,10 +522,6 @@ pub fn shader_defs( #[cfg(feature = "morph_particles")] shader_defs.push("READ_WRITE_POINTS".into()); - if key.visualize_depth { - shader_defs.push("VISUALIZE_DEPTH".into()); - } - #[cfg(feature = "packed")] shader_defs.push("PACKED".into()); @@ -561,6 +558,12 @@ pub fn shader_defs( #[cfg(feature = "webgl2")] shader_defs.push("WEBGL2".into()); + match key.rasterize_mode { + GaussianCloudRasterize::Color => {}, + GaussianCloudRasterize::Depth => shader_defs.push("RASTERIZE_DEPTH".into()), + GaussianCloudRasterize::Normal => shader_defs.push("RASTERIZE_NORMAL".into()), + } + match key.draw_mode { GaussianCloudDrawMode::All => {}, GaussianCloudDrawMode::Selected => shader_defs.push("DRAW_SELECTED".into()), @@ -574,8 +577,8 @@ pub fn shader_defs( pub struct GaussianCloudPipelineKey { pub aabb: bool, pub visualize_bounding_box: bool, - pub visualize_depth: bool, pub draw_mode: GaussianCloudDrawMode, + pub rasterize_mode: GaussianCloudRasterize, pub sample_count: u32, }