diff --git a/examples/multi_camera.rs b/examples/multi_camera.rs index ff6b33da..450f855e 100644 --- a/examples/multi_camera.rs +++ b/examples/multi_camera.rs @@ -6,7 +6,7 @@ use bevy_inspector_egui::{bevy_egui::EguiPlugin, quick::WorldInspectorPlugin}; use bevy_panorbit_camera::{PanOrbitCamera, PanOrbitCameraPlugin}; use bevy_gaussian_splatting::{ - CloudSettings, Gaussian3d, GaussianCamera, GaussianMode, GaussianSplattingPlugin, + CloudSettings, Gaussian3d, GaussianBounds, GaussianCamera, GaussianMode, GaussianSplattingPlugin, PlanarGaussian3d, PlanarGaussian3dHandle, SphericalHarmonicCoefficients, gaussian::f32::Rotation, utils::{GaussianSplattingViewer, setup_hooks}, @@ -139,7 +139,7 @@ pub fn setup_multi_camera( Transform::from_translation(Vec3::new(spacing, spacing, 0.0)), PlanarGaussian3dHandle(gaussian_assets.add(red_gaussians)), CloudSettings { - aabb: true, + bounds: GaussianBounds::Aabb, gaussian_mode: GaussianMode::Gaussian2d, ..default() }, diff --git a/src/gaussian/settings.rs b/src/gaussian/settings.rs index d5d22d9a..6977c9b5 100644 --- a/src/gaussian/settings.rs +++ b/src/gaussian/settings.rs @@ -32,6 +32,25 @@ pub enum PlaybackMode { Still, } +#[derive( + Clone, Copy, Debug, Default, Eq, Hash, PartialEq, Reflect, Serialize, Deserialize, ValueEnum, +)] +pub enum GaussianBounds { + Aabb, + #[default] + Obb, + Triangle, +} + +impl GaussianBounds { + pub const fn vertex_count(self) -> u32 { + match self { + GaussianBounds::Triangle => 3, + GaussianBounds::Aabb | GaussianBounds::Obb => 4, + } + } +} + #[derive( Clone, Copy, Debug, Default, Eq, Hash, PartialEq, Reflect, Serialize, Deserialize, ValueEnum, )] @@ -51,7 +70,7 @@ pub enum RasterizeMode { #[reflect(Component)] #[serde(default)] pub struct CloudSettings { - pub aabb: bool, + pub bounds: GaussianBounds, pub global_opacity: f32, pub global_scale: f32, pub opacity_adaptive_radius: bool, @@ -71,7 +90,7 @@ pub struct CloudSettings { impl Default for CloudSettings { fn default() -> Self { Self { - aabb: false, + bounds: GaussianBounds::default(), global_opacity: 1.0, global_scale: 1.0, opacity_adaptive_radius: true, @@ -94,6 +113,7 @@ impl Default for CloudSettings { pub struct SettingsPlugin; impl Plugin for SettingsPlugin { fn build(&self, app: &mut App) { + app.register_type::(); app.register_type::(); app.add_systems(Update, (playback_update,)); diff --git a/src/lib.rs b/src/lib.rs index 0a1d716c..f26e3a61 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,7 +11,7 @@ pub use gaussian::{ planar_3d::{Gaussian3d, PlanarGaussian3d, PlanarGaussian3dHandle, random_gaussians_3d}, planar_4d::{Gaussian4d, PlanarGaussian4d, PlanarGaussian4dHandle, random_gaussians_4d}, }, - settings::{CloudSettings, GaussianMode, RasterizeMode}, + settings::{CloudSettings, GaussianBounds, GaussianMode, RasterizeMode}, }; pub use io::scene::{GaussianScene, GaussianSceneHandle}; diff --git a/src/render/bindings.wgsl b/src/render/bindings.wgsl index 481caf25..b1aa5849 100644 --- a/src/render/bindings.wgsl +++ b/src/render/bindings.wgsl @@ -16,6 +16,7 @@ struct GaussianUniforms { global_scale: f32, count: u32, count_root_ceil: u32, + vertex_count: u32, time: f32, time_start: f32, time_stop: f32, diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index 1e80284d..1900c8a0 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -216,15 +216,28 @@ fn vs_points( return output; } - var quad_vertices = array, 4>( - vec2(-1.0, -1.0), - vec2(-1.0, 1.0), - vec2( 1.0, -1.0), - vec2( 1.0, 1.0), - ); + #ifdef USE_TRIANGLE + let vertex_count: u32 = 3u; + let sqrt3 = sqrt(3.0); + var bounds_vertices = array, 3>( + vec2(0.0, 2.0), + vec2( sqrt3, -1.0), + vec2(-sqrt3, -1.0), + ); + #else + let vertex_count: u32 = 4u; + var bounds_vertices = array, 4>( + vec2(-1.0, -1.0), + vec2(-1.0, 1.0), + vec2( 1.0, -1.0), + vec2( 1.0, 1.0), + ); + #endif + + let bounds_index = vertex_index % vertex_count; + let bounds_offset = bounds_vertices[bounds_index]; - let quad_index = vertex_index % 4u; - let quad_offset = quad_vertices[quad_index]; + let vertex_uv = bounds_offset; var opacity = get_opacity(splat_index); @@ -248,7 +261,7 @@ fn vs_points( let bb = get_bounding_box_cov2d( surfel.extent, - quad_offset, + bounds_offset, cutoff, ); output.radius = bb.zw; @@ -292,7 +305,7 @@ fn vs_points( let bb = get_bounding_box_clip( gaussian_cov2d, - quad_offset, + bounds_offset, cutoff, ); @@ -307,6 +320,18 @@ fn vs_points( output.conic = conic; output.major_minor = bb.zw; #endif + + #ifdef USE_TRIANGLE + let det = gaussian_cov2d.x * gaussian_cov2d.z - gaussian_cov2d.y * gaussian_cov2d.y; + let det_inv = 1.0 / det; + let conic = vec3( + gaussian_cov2d.z * det_inv, + -gaussian_cov2d.y * det_inv, + gaussian_cov2d.x * det_inv + ); + output.conic = conic; + output.major_minor = bb.zw; + #endif #endif var rgb = vec3(0.0); @@ -426,7 +451,7 @@ fn vs_points( } #endif - output.uv = quad_offset; + output.uv = vertex_uv; output.position = vec4( projected_position.xy + bb.xy, projected_position.zw, @@ -437,35 +462,71 @@ fn vs_points( @fragment fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { + var power: f32 = 0.0; + #ifdef USE_AABB -#ifdef GAUSSIAN_2D - let radius = input.radius; - let mean_2d = input.mean_2d; - let aspect = vec2( - 1.0, - view.viewport.z / view.viewport.w, - ); - let pixel_coord = input.uv * radius * aspect + mean_2d; - - let power = surfel_fragment_power( - mat3x3( - input.local_to_pixel_u, - input.local_to_pixel_v, - input.local_to_pixel_w, - ), - pixel_coord, - mean_2d, - ); -#else ifdef GAUSSIAN_3D - let d = -input.major_minor; - let conic = input.conic; - let power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y; -#else ifdef GAUSSIAN_4D - let d = -input.major_minor; - let conic = input.conic; - let power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y; + #ifdef GAUSSIAN_2D + let radius = input.radius; + let mean_2d = input.mean_2d; + let aspect = vec2( + 1.0, + view.viewport.z / view.viewport.w, + ); + let pixel_coord = input.uv * radius * aspect + mean_2d; + + power = surfel_fragment_power( + mat3x3( + input.local_to_pixel_u, + input.local_to_pixel_v, + input.local_to_pixel_w, + ), + pixel_coord, + mean_2d, + ); + #else ifdef GAUSSIAN_3D + let d = -input.major_minor; + let conic = input.conic; + power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y; + #else ifdef GAUSSIAN_4D + let d = -input.major_minor; + let conic = input.conic; + power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y; + #endif + + if (power > 0.0) { + discard; + } #endif +#ifdef USE_TRIANGLE + #ifdef GAUSSIAN_2D + let radius = input.radius; + let mean_2d = input.mean_2d; + let aspect = vec2( + 1.0, + view.viewport.z / view.viewport.w, + ); + let pixel_coord = input.uv * radius * aspect + mean_2d; + + power = surfel_fragment_power( + mat3x3( + input.local_to_pixel_u, + input.local_to_pixel_v, + input.local_to_pixel_w, + ), + pixel_coord, + mean_2d, + ); + #else ifdef GAUSSIAN_3D + let d = -input.major_minor; + let conic = input.conic; + power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y; + #else ifdef GAUSSIAN_4D + let d = -input.major_minor; + let conic = input.conic; + power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y; + #endif + if (power > 0.0) { discard; } @@ -476,7 +537,7 @@ fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { let sigma_squared = 2.0 * sigma * sigma; let distance_squared = dot(input.uv, input.uv); - let power = -distance_squared / sigma_squared; + power = -distance_squared / sigma_squared; if (distance_squared > 3.0 * 3.0) { discard; @@ -484,14 +545,30 @@ fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { #endif #ifdef VISUALIZE_BOUNDING_BOX - let uv = input.uv * 0.5 + 0.5; - let edge_width = 0.08; - if ( - (uv.x < edge_width || uv.x > 1.0 - edge_width) || - (uv.y < edge_width || uv.y > 1.0 - edge_width) - ) { - return vec4(0.3, 1.0, 0.1, 1.0); - } + #ifdef USE_TRIANGLE + let sqrt3 = sqrt(3.0); + let v0 = vec2(0.0, 2.0); + let v1 = vec2( sqrt3, -1.0); + let v2 = vec2(-sqrt3, -1.0); + let denom = (v1.y - v2.y) * (v0.x - v2.x) + (v2.x - v1.x) * (v0.y - v2.y); + let bary0 = ((v1.y - v2.y) * (input.uv.x - v2.x) + (v2.x - v1.x) * (input.uv.y - v2.y)) / denom; + let bary1 = ((v2.y - v0.y) * (input.uv.x - v2.x) + (v0.x - v2.x) * (input.uv.y - v2.y)) / denom; + let bary2 = 1.0 - bary0 - bary1; + let edge_width = 0.05; + let min_bary = min(min(bary0, bary1), bary2); + if (min_bary < edge_width) { + return vec4(0.3, 1.0, 0.1, 1.0); + } + #else + let uv = input.uv * 0.5 + 0.5; + let edge_width = 0.08; + if ( + (uv.x < edge_width || uv.x > 1.0 - edge_width) || + (uv.y < edge_width || uv.y > 1.0 - edge_width) + ) { + return vec4(0.3, 1.0, 0.1, 1.0); + } + #endif #endif let alpha = min(exp(power) * input.color.a, 0.999); diff --git a/src/render/helpers.wgsl b/src/render/helpers.wgsl index 0fb1fc20..18e84587 100644 --- a/src/render/helpers.wgsl +++ b/src/render/helpers.wgsl @@ -78,6 +78,21 @@ fn get_bounding_box_clip( ); #endif +#ifdef USE_TRIANGLE + let radius_px = cutoff * max(x_axis_length, y_axis_length); + let radius_ndc = vec2( + radius_px / view.viewport.z, + radius_px / view.viewport.w, + ); + let vertex_px = direction * radius_px; + let ndc_vertex = direction * radius_ndc; + + return vec4( + ndc_vertex, + vertex_px, + ); +#endif + #ifdef USE_OBB let a = (cov2d.x - cov2d.z) * (cov2d.x - cov2d.z); diff --git a/src/render/mod.rs b/src/render/mod.rs index cfc49610..2f40c53e 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -28,7 +28,7 @@ use crate::{ gaussian::{ cloud::CloudVisibilityClass, interface::CommonCloud, - settings::{CloudSettings, DrawMode, GaussianMode, RasterizeMode}, + settings::{CloudSettings, DrawMode, GaussianBounds, GaussianMode, RasterizeMode}, }, material::{ spherical_harmonics::{HALF_SH_COEFF_COUNT, SH_COEFF_COUNT, SH_DEGREE, SH_VEC4_PLANES}, @@ -407,7 +407,7 @@ fn queue_gaussians( let msaa = msaa.cloned().unwrap_or_default(); let key = CloudPipelineKey { - aabb: settings.aabb, + bounds: settings.bounds, binary_gaussian_op: false, opacity_adaptive_radius: settings.opacity_adaptive_radius, visualize_bounding_box: settings.visualize_bounding_box, @@ -714,12 +714,10 @@ pub fn shader_defs(key: CloudPipelineKey) -> Vec { ), ]; - if key.aabb { - shader_defs.push("USE_AABB".into()); - } - - if !key.aabb { - shader_defs.push("USE_OBB".into()); + match key.bounds { + GaussianBounds::Aabb => shader_defs.push("USE_AABB".into()), + GaussianBounds::Obb => shader_defs.push("USE_OBB".into()), + GaussianBounds::Triangle => shader_defs.push("USE_TRIANGLE".into()), } if key.binary_gaussian_op { @@ -806,7 +804,7 @@ pub fn shader_defs(key: CloudPipelineKey) -> Vec { #[derive(PartialEq, Eq, Hash, Clone, Copy, Default)] pub struct CloudPipelineKey { - pub aabb: bool, + pub bounds: GaussianBounds, pub binary_gaussian_op: bool, pub visualize_bounding_box: bool, pub opacity_adaptive_radius: bool, @@ -907,6 +905,7 @@ pub struct CloudUniform { pub global_scale: f32, pub count: u32, pub count_root_ceil: u32, + pub vertex_count: u32, pub time: f32, pub time_start: f32, pub time_stop: f32, @@ -966,6 +965,7 @@ pub fn extract_gaussians( global_scale: settings.global_scale, count: cloud.len() as u32, count_root_ceil: (cloud.len() as f32).sqrt().ceil() as u32, + vertex_count: settings.bounds.vertex_count(), time: settings.time, time_start: settings.time_start, time_stop: settings.time_stop, diff --git a/src/sort/radix.wgsl b/src/sort/radix.wgsl index d15d2a29..439efaf7 100644 --- a/src/sort/radix.wgsl +++ b/src/sort/radix.wgsl @@ -75,7 +75,8 @@ fn radix_sort_a( @builtin(global_invocation_id) gl_GlobalInvocationID: vec3, ) { if (gl_LocalInvocationID.x == 0u && gl_LocalInvocationID.y == 0u && gl_GlobalInvocationID.x == 0u) { - draw_indirect.vertex_count = 4u; + let vertex_count = max(gaussian_uniforms.vertex_count, 3u); + draw_indirect.vertex_count = vertex_count; atomicStore(&draw_indirect.instance_count, gaussian_uniforms.count); } workgroupBarrier(); diff --git a/tools/compare_aabb_obb.rs b/tools/compare_aabb_obb.rs index 0e08a623..00a416bb 100644 --- a/tools/compare_aabb_obb.rs +++ b/tools/compare_aabb_obb.rs @@ -5,7 +5,7 @@ use bevy_interleave::prelude::Planar; use bevy_panorbit_camera::{PanOrbitCamera, PanOrbitCameraPlugin}; use bevy_gaussian_splatting::{ - CloudSettings, Gaussian3d, GaussianCamera, GaussianSplattingPlugin, PlanarGaussian3d, + CloudSettings, Gaussian3d, GaussianBounds, GaussianCamera, GaussianSplattingPlugin, PlanarGaussian3d, PlanarGaussian3dHandle, SphericalHarmonicCoefficients, utils::{GaussianSplattingViewer, setup_hooks}, }; @@ -30,7 +30,7 @@ pub fn setup_aabb_obb_compare( blue_aabb_gaussian, ]))), CloudSettings { - aabb: true, + bounds: GaussianBounds::Aabb, visualize_bounding_box: true, ..default() }, @@ -53,7 +53,7 @@ pub fn setup_aabb_obb_compare( red_obb_gaussian, ]))), CloudSettings { - aabb: false, + bounds: GaussianBounds::Obb, visualize_bounding_box: true, ..default() }, diff --git a/tools/surfel_plane.rs b/tools/surfel_plane.rs index 1cf89ce6..75819d45 100644 --- a/tools/surfel_plane.rs +++ b/tools/surfel_plane.rs @@ -5,7 +5,7 @@ use bevy_interleave::prelude::Planar; use bevy_panorbit_camera::{PanOrbitCamera, PanOrbitCameraPlugin}; use bevy_gaussian_splatting::{ - CloudSettings, Gaussian3d, GaussianCamera, GaussianMode, GaussianSplattingPlugin, + CloudSettings, Gaussian3d, GaussianBounds, GaussianCamera, GaussianMode, GaussianSplattingPlugin, PlanarGaussian3d, PlanarGaussian3dHandle, SphericalHarmonicCoefficients, gaussian::f32::Rotation, utils::{GaussianSplattingViewer, setup_hooks}, @@ -96,7 +96,7 @@ pub fn setup_surfel_compare( PlanarGaussian3dHandle(cloud), CloudSettings { visualize_bounding_box, - aabb: true, + bounds: GaussianBounds::Aabb, gaussian_mode: GaussianMode::Gaussian2d, ..default() },