diff --git a/Cargo.toml b/Cargo.toml index 88ad6ea4..c4161fa7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "bevy_gaussian_splatting" description = "bevy gaussian splatting render pipeline plugin" -version = "2.5.0" +version = "2.6.0" edition = "2021" authors = ["mosure "] license = "MIT" @@ -234,6 +234,11 @@ name = "compare_aabb_obb" path = "tools/compare_aabb_obb.rs" required-features = ["tooling"] +[[bin]] +name = "surfel_plane" +path = "tools/surfel_plane.rs" +required-features = ["tooling"] + [[bin]] name = "test_gaussian" path = "tests/gpu/gaussian.rs" diff --git a/README.md b/README.md index 110f2e30..5992de43 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,8 @@ bevy gaussian splatting render pipeline plugin. view the [live demo](https://mos - [X] depth colorization - [X] f16 and f32 gcloud - [X] wgl2 and webgpu +- [X] 3dgs +- [X] 2dgs - [ ] spherical harmonic coefficients clustering - [ ] 4D gaussian cloud wavelet compression - [ ] accelerated spatial queries @@ -83,6 +85,7 @@ fn setup_gaussian_cloud( # credits +- [2d-gaussian-splatting](https://github.com/hbb1/2d-gaussian-splatting) - [4d gaussians](https://github.com/hustvl/4DGaussians) - [4d-gaussian-splatting](https://fudan-zvg.github.io/4d-gaussian-splatting/) - [bevy](https://github.com/bevyengine/bevy) diff --git a/src/gaussian/settings.rs b/src/gaussian/settings.rs index 51ebce2e..2215b278 100644 --- a/src/gaussian/settings.rs +++ b/src/gaussian/settings.rs @@ -21,6 +21,23 @@ pub enum GaussianCloudDrawMode { } +#[derive( + Clone, + Copy, + Debug, + Default, + Eq, + Hash, + PartialEq, + Reflect, +)] +pub enum GaussianMode { + #[default] + Gaussian3d, + GaussianSurfel, +} + + #[derive( Clone, Copy, @@ -46,9 +63,11 @@ pub struct GaussianCloudSettings { pub global_opacity: f32, pub global_scale: f32, pub transform: Transform, + pub opacity_adaptive_radius: bool, pub visualize_bounding_box: bool, pub sort_mode: SortMode, pub draw_mode: GaussianCloudDrawMode, + pub gaussian_mode: GaussianMode, pub rasterize_mode: GaussianCloudRasterize, } @@ -59,9 +78,11 @@ impl Default for GaussianCloudSettings { global_opacity: 1.0, global_scale: 1.0, transform: Transform::IDENTITY, + opacity_adaptive_radius: true, visualize_bounding_box: false, sort_mode: SortMode::default(), draw_mode: GaussianCloudDrawMode::default(), + gaussian_mode: GaussianMode::default(), rasterize_mode: GaussianCloudRasterize::default(), } } diff --git a/src/lib.rs b/src/lib.rs index dc4a802c..a0c071a0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,7 +4,11 @@ pub use gaussian::{ packed::Gaussian, cloud::GaussianCloud, rand::random_gaussians, - settings::GaussianCloudSettings, + settings::{ + GaussianCloudRasterize, + GaussianCloudSettings, + GaussianMode, + }, }; pub use material::spherical_harmonics::SphericalHarmonicCoefficients; diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index bbafd7ff..3d4c271f 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -1,6 +1,5 @@ #import bevy_gaussian_splatting::bindings::{ view, - globals, gaussian_uniforms, sorting_pass_index, sorting, @@ -12,6 +11,15 @@ #import bevy_gaussian_splatting::depth::{ depth_to_rgb, } +#import bevy_gaussian_splatting::helpers::{ + get_rotation_matrix, + get_scale_matrix, +} +#import bevy_gaussian_splatting::surfel::{ + compute_cov2d_surfel, + get_bounding_box_cov2d, + surfel_fragment_power, +} #import bevy_gaussian_splatting::transform::{ world_to_clip, in_frustum, @@ -115,55 +123,37 @@ fn get_entry(index: u32) -> Entry { struct GaussianVertexOutput { @builtin(position) position: vec4, @location(0) color: vec4, - @location(1) conic: vec3, - @location(2) uv: vec2, + @location(1) uv: vec2, +#ifdef GAUSSIAN_3D + @location(2) conic: vec3, @location(3) major_minor: vec2, +#else ifdef GAUSSIAN_SURFEL + @location(2) local_to_pixel_u: vec3, + @location(3) local_to_pixel_v: vec3, + @location(4) local_to_pixel_w: vec3, + @location(5) mean_2d: vec2, + @location(6) radius: vec2, +#endif }; #else struct GaussianVertexOutput { @builtin(position) position: vec4, @location(0) @interpolate(flat) color: vec4, - @location(1) @interpolate(flat) conic: vec3, - @location(2) @interpolate(linear) uv: vec2, + @location(1) @interpolate(linear) uv: vec2, +#ifdef GAUSSIAN_3D + @location(2) @interpolate(flat) conic: vec3, @location(3) @interpolate(linear) major_minor: vec2, +#else ifdef GAUSSIAN_SURFEL + @location(2) @interpolate(flat) local_to_pixel_u: vec3, + @location(3) @interpolate(flat) local_to_pixel_v: vec3, + @location(4) @interpolate(flat) local_to_pixel_w: vec3, + @location(5) @interpolate(flat) mean_2d: vec2, + @location(6) @interpolate(flat) radius: vec2, +#endif }; #endif -fn get_rotation_matrix( - rotation: vec4, -) -> mat3x3 { - let r = rotation.x; - let x = rotation.y; - let y = rotation.z; - let z = rotation.w; - - return mat3x3( - 1.0 - 2.0 * (y * y + z * z), - 2.0 * (x * y - r * z), - 2.0 * (x * z + r * y), - - 2.0 * (x * y + r * z), - 1.0 - 2.0 * (x * x + z * z), - 2.0 * (y * z - r * x), - - 2.0 * (x * z - r * y), - 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 { @@ -191,7 +181,7 @@ fn compute_cov3d(scale: vec3, rotation: vec4) -> array { ); } -fn compute_cov2d( +fn compute_cov2d_3dgs( position: vec3, index: u32, ) -> vec3 { @@ -213,8 +203,8 @@ fn compute_cov2d( var t = view.view_from_world * vec4(position, 1.0); let focal = vec2( - view.clip_from_view .x.x * view.viewport.z, - view.clip_from_view .y.y * view.viewport.w, + view.clip_from_view.x.x * view.viewport.z, + view.clip_from_view.y.y * view.viewport.w, ); let s = 1.0 / (t.z * t.z); @@ -244,6 +234,7 @@ fn compute_cov2d( fn get_bounding_box( cov2d: vec3, direction: vec2, + cutoff: f32, ) -> vec4 { // return vec4(offset, uv); @@ -262,7 +253,7 @@ fn get_bounding_box( #ifdef USE_AABB - let radius_px = 3.5 * max(x_axis_length, y_axis_length); + let radius_px = cutoff * max(x_axis_length, y_axis_length); let radius_ndc = vec2( radius_px / view.viewport.zw, ); @@ -280,7 +271,7 @@ fn get_bounding_box( let major_radius = sqrt((cov2d.x + cov2d.z + b) * 0.5); let minor_radius = sqrt((cov2d.x + cov2d.z - b) * 0.5); - let bounds = 3.5 * vec2( + let bounds = cutoff * vec2( major_radius, minor_radius, ); @@ -400,28 +391,19 @@ fn vs_points( max_distance, ); #else ifdef RASTERIZE_NORMAL + let R = get_rotation_matrix(get_rotation(splat_index)); + let S = get_scale_matrix(get_scale(splat_index)); let T = mat3x3( gaussian_uniforms.transform[0].xyz, gaussian_uniforms.transform[1].xyz, gaussian_uniforms.transform[2].xyz, ); + let L = T * S * R; - let R = get_rotation_matrix(get_rotation(splat_index)); - let scale = get_scale(splat_index); - let scale_inf = inverted_infinity_norm(scale); - let S = get_scale_matrix(scale_inf); + let local_normal = vec4(L.z, 0.0); + let world_normal = view.view_from_world * local_normal; - 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); + let t = normalize(world_normal); rgb = vec3( 0.5 * (t.x + 1.0), @@ -432,10 +414,18 @@ fn vs_points( rgb = get_color(splat_index, ray_direction); #endif + let opacity = get_opacity(splat_index); + +#ifdef OPACITY_ADAPTIVE_RADIUS + let cutoff = sqrt(max(9.0 + 2.0 * log(opacity), 0.000001)); +#else + let cutoff = 3.0; +#endif + // TODO: verify color benefit for ray_direction computed at quad verticies instead of gaussian center (same as current complexity) output.color = vec4( rgb, - get_opacity(splat_index) * gaussian_uniforms.global_opacity, + opacity * gaussian_uniforms.global_opacity, ); #ifdef HIGHLIGHT_SELECTED @@ -444,7 +434,16 @@ fn vs_points( } #endif - let cov2d = compute_cov2d(transformed_position, splat_index); +#ifdef GAUSSIAN_3D + let cov2d = compute_cov2d_3dgs( + transformed_position, + splat_index, + ); + let bb = get_bounding_box( + cov2d, + quad_offset, + cutoff, + ); #ifdef USE_AABB let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y; @@ -454,19 +453,35 @@ fn vs_points( -cov2d.y * det_inv, cov2d.x * det_inv ); + // TODO: this conic seems only valid in 3dgs output.conic = conic; + output.major_minor = bb.zw; #endif - let bb = get_bounding_box( - cov2d, +#else ifdef GAUSSIAN_SURFEL + let surfel = compute_cov2d_surfel( + transformed_position, + splat_index, + cutoff, + ); + + output.local_to_pixel_u = surfel.local_to_pixel.x; + output.local_to_pixel_v = surfel.local_to_pixel.y; + output.local_to_pixel_w = surfel.local_to_pixel.z; + output.mean_2d = surfel.mean_2d; + + let bb = get_bounding_box_cov2d( + surfel.extent, quad_offset, + cutoff, ); + output.radius = bb.zw; +#endif output.uv = quad_offset; - output.major_minor = bb.zw; output.position = vec4( projected_position.xy + bb.xy, - projected_position.zw + projected_position.zw, ); return output; @@ -475,9 +490,30 @@ fn vs_points( @fragment fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { #ifdef USE_AABB +#ifdef GAUSSIAN_SURFEL + 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 pixel_coord = input.position.xy * view.viewport.zw + view.viewport.xy; + + 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; +#endif if (power > 0.0) { discard; @@ -485,19 +521,19 @@ fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { #endif #ifdef USE_OBB - let sigma = 1.0 / 3.5; + let sigma = 1.0 / 3.0; let sigma_squared = 2.0 * sigma * sigma; let distance_squared = dot(input.uv, input.uv); let power = -distance_squared / sigma_squared; - if (distance_squared > 3.5 * 3.5) { + if (distance_squared > 3.0 * 3.0) { discard; } #endif #ifdef VISUALIZE_BOUNDING_BOX - let uv = (input.uv + 1.0) / 2.0; + let uv = input.uv * 0.5 + 0.5; let edge_width = 0.08; if ( (uv.x < edge_width || uv.x > 1.0 - edge_width) || @@ -507,13 +543,12 @@ fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { } #endif - let alpha = exp(power); - let final_alpha = alpha * input.color.a; + let alpha = min(exp(power) * input.color.a, 0.999); - // TODO: round final_alpha to terminate depth test? + // TODO: round alpha to terminate depth test? return vec4( - input.color.rgb * final_alpha, - final_alpha, + input.color.rgb * alpha, + alpha, ); } diff --git a/src/render/helpers.wgsl b/src/render/helpers.wgsl new file mode 100644 index 00000000..60989364 --- /dev/null +++ b/src/render/helpers.wgsl @@ -0,0 +1,57 @@ +#define_import_path bevy_gaussian_splatting::helpers + +#import bevy_gaussian_splatting::bindings::{ + view, + gaussian_uniforms, +} + + +fn intrinsic_matrix() -> mat3x4 { + let focal = vec2( + view.clip_from_view.x.x * view.viewport.z / 2.0, + view.clip_from_view.y.y * view.viewport.w / 2.0, + ); + + let Ks = mat3x4( + vec4(focal.x, 0.0, 0.0, (view.viewport.z - 1.0) / 2.0), + vec4(0.0, focal.y, 0.0, (view.viewport.w - 1.0) / 2.0), + vec4(0.0, 0.0, 0.0, 1.0) + ); + + return Ks; +} + + +fn get_rotation_matrix( + rotation: vec4, +) -> mat3x3 { + let r = rotation.x; + let x = rotation.y; + let y = rotation.z; + let z = rotation.w; + + return mat3x3( + 1.0 - 2.0 * (y * y + z * z), + 2.0 * (x * y - r * z), + 2.0 * (x * z + r * y), + + 2.0 * (x * y + r * z), + 1.0 - 2.0 * (x * x + z * z), + 2.0 * (y * z - r * x), + + 2.0 * (x * z - r * y), + 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, + ); +} diff --git a/src/render/mod.rs b/src/render/mod.rs index b0083780..4715bc59 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -66,6 +66,7 @@ use crate::{ GaussianCloudDrawMode, GaussianCloudRasterize, GaussianCloudSettings, + GaussianMode, }, }, material::spherical_harmonics::{ @@ -93,6 +94,8 @@ mod texture; const BINDINGS_SHADER_HANDLE: Handle = Handle::weak_from_u128(675257236); const GAUSSIAN_SHADER_HANDLE: Handle = Handle::weak_from_u128(68294581); +const GAUSSIAN_SURFEL_SHADER_HANDLE: Handle = Handle::weak_from_u128(123166726); +const HELPERS_SHADER_HANDLE: Handle = Handle::weak_from_u128(134646367); const PACKED_SHADER_HANDLE: Handle = Handle::weak_from_u128(123623514); const PLANAR_SHADER_HANDLE: Handle = Handle::weak_from_u128(72345231); const TEXTURE_SHADER_HANDLE: Handle = Handle::weak_from_u128(26345735); @@ -118,6 +121,20 @@ impl Plugin for RenderPipelinePlugin { Shader::from_wgsl ); + load_internal_asset!( + app, + GAUSSIAN_SURFEL_SHADER_HANDLE, + "surfel.wgsl", + Shader::from_wgsl + ); + + load_internal_asset!( + app, + HELPERS_SHADER_HANDLE, + "helpers.wgsl", + Shader::from_wgsl + ); + load_internal_asset!( app, PACKED_SHADER_HANDLE, @@ -315,8 +332,10 @@ fn queue_gaussians( let key = GaussianCloudPipelineKey { aabb: settings.aabb, + opacity_adaptive_radius: settings.opacity_adaptive_radius, visualize_bounding_box: settings.visualize_bounding_box, draw_mode: settings.draw_mode, + gaussian_mode: settings.gaussian_mode, rasterize_mode: settings.rasterize_mode, sample_count: msaa.samples(), }; @@ -525,6 +544,10 @@ pub fn shader_defs( shader_defs.push("USE_OBB".into()); } + if key.opacity_adaptive_radius { + shader_defs.push("OPACITY_ADAPTIVE_RADIUS".into()); + } + if key.visualize_bounding_box { shader_defs.push("VISUALIZE_BOUNDING_BOX".into()); } @@ -568,8 +591,13 @@ pub fn shader_defs( #[cfg(feature = "webgl2")] shader_defs.push("WEBGL2".into()); + match key.gaussian_mode { + GaussianMode::Gaussian3d => shader_defs.push("GAUSSIAN_3D".into()), + GaussianMode::GaussianSurfel => shader_defs.push("GAUSSIAN_SURFEL".into()), + } + match key.rasterize_mode { - GaussianCloudRasterize::Color => {}, + GaussianCloudRasterize::Color => shader_defs.push("RASTERIZE_COLOR".into()), GaussianCloudRasterize::Depth => shader_defs.push("RASTERIZE_DEPTH".into()), GaussianCloudRasterize::Normal => shader_defs.push("RASTERIZE_NORMAL".into()), } @@ -587,7 +615,9 @@ pub fn shader_defs( pub struct GaussianCloudPipelineKey { pub aabb: bool, pub visualize_bounding_box: bool, + pub opacity_adaptive_radius: bool, pub draw_mode: GaussianCloudDrawMode, + pub gaussian_mode: GaussianMode, pub rasterize_mode: GaussianCloudRasterize, pub sample_count: u32, } diff --git a/src/render/surfel.wgsl b/src/render/surfel.wgsl new file mode 100644 index 00000000..7d9ab1e5 --- /dev/null +++ b/src/render/surfel.wgsl @@ -0,0 +1,197 @@ +#define_import_path bevy_gaussian_splatting::surfel + +#import bevy_gaussian_splatting::bindings::{ + view, + gaussian_uniforms, +} +#import bevy_gaussian_splatting::helpers::{ + get_rotation_matrix, + get_scale_matrix, + intrinsic_matrix, +} + +#ifdef PACKED +#ifdef PRECOMPUTE_COVARIANCE_3D +#import bevy_gaussian_splatting::packed::{ + get_position, + get_color, + get_visibility, + get_opacity, + get_cov3d, +} +#else +#import bevy_gaussian_splatting::packed::{ + get_position, + get_color, + get_visibility, + get_opacity, + get_rotation, + get_scale, +} +#endif +#else + +#ifdef BUFFER_STORAGE +#ifdef PRECOMPUTE_COVARIANCE_3D +#import bevy_gaussian_splatting::planar::{ + get_position, + get_color, + get_visibility, + get_opacity, + get_cov3d, +} +#else +#import bevy_gaussian_splatting::planar::{ + get_position, + get_color, + get_visibility, + get_opacity, + get_rotation, + get_scale, +} +#endif +#endif + +#endif + + +#ifdef BUFFER_TEXTURE +#ifdef PRECOMPUTE_COVARIANCE_3D +#import bevy_gaussian_splatting::texture::{ + get_position, + get_color, + get_visibility, + get_opacity, + get_cov3d, + location, +} +#else +#import bevy_gaussian_splatting::texture::{ + get_position, + get_color, + get_visibility, + get_opacity, + get_rotation, + get_scale, + location, +} +#endif +#endif + + +struct Surfel { + local_to_pixel: mat3x3, + mean_2d: vec2, + extent: vec2, +}; + + +fn get_bounding_box_cov2d( + extent: vec2, + direction: vec2, + cutoff: f32, +) -> vec4 { + let fitler_size = 0.707106; + + if extent.x < 1.e-4 || extent.y < 1.e-4 { + return vec4(0.0); + } + + let radius = sqrt(extent); + let max_radius = vec2(max( + max(radius.x, radius.y), + cutoff * fitler_size, + )); + + // TODO: verify OBB capability + let radius_ndc = vec2( + max_radius / view.viewport.zw, + ); + + return vec4( + radius_ndc * direction, + max_radius, + ); +} + + +fn compute_cov2d_surfel( + gaussian_position: vec3, + index: u32, + cutoff: f32, +) -> Surfel { + var output: Surfel; + + let rotation = get_rotation(index); + let scale = get_scale(index); + + let T_r = mat3x3( + gaussian_uniforms.transform[0].xyz, + gaussian_uniforms.transform[1].xyz, + gaussian_uniforms.transform[2].xyz, + ); + + let S = get_scale_matrix(scale); + let R = get_rotation_matrix(rotation); + + let L = T_r * S * R; + + let world_from_local = mat3x4( + vec4(L.x, 0.0), + vec4(L.y, 0.0), + vec4(gaussian_position, 1.0), + ); + + let ndc_from_world = transpose(view.clip_from_world); + let pixels_from_ndc = intrinsic_matrix(); + + let T = transpose(world_from_local) * ndc_from_world * pixels_from_ndc; + + let test = vec3(cutoff * cutoff, cutoff * cutoff, -1.0); + let d = dot(test * T[2], T[2]); + if abs(d) < 1.0e-4 { + output.extent = vec2(0.0); + return output; + } + + let f = (1.0 / d) * test; + let mean_2d = vec2( + dot(f, T[0] * T[2]), + dot(f, T[1] * T[2]), + ); + + let t = vec2( + dot(f * T[0], T[0]), + dot(f * T[1], T[1]), + ); + let extent = mean_2d * mean_2d - t; + + output.local_to_pixel = T; + output.mean_2d = mean_2d; + output.extent = extent; + return output; +} + +fn surfel_fragment_power( + local_to_pixel: mat3x3, + pixel_coord: vec2, + mean_2d: vec2, +) -> f32 { + let deltas = mean_2d - pixel_coord; + + let hu = pixel_coord.x * local_to_pixel.z - local_to_pixel.x; + let hv = pixel_coord.y * local_to_pixel.z - local_to_pixel.y; + + let p = cross(hu, hv); + + let us = p.x / p.z; + let vs = p.y / p.z; + + let sigmas_3d = us * us + vs * vs; + let sigmas_2d = 2.0 * (deltas.x * deltas.x + deltas.y * deltas.y); + + let sigmas = 0.5 * min(sigmas_3d, sigmas_2d); + let power = -sigmas; + + return power; +} diff --git a/src/sort/rayon.rs b/src/sort/rayon.rs index a1f98190..b64201a4 100644 --- a/src/sort/rayon.rs +++ b/src/sort/rayon.rs @@ -115,7 +115,7 @@ pub fn rayon_sort( }); sorted_entries.sorted.par_sort_unstable_by(|a, b| { - bytemuck::cast::(b.key).partial_cmp(&bytemuck::cast::(a.key)).unwrap() + bytemuck::cast::(b.key).partial_cmp(&bytemuck::cast::(a.key)).unwrap_or(std::cmp::Ordering::Equal) }); // TODO: update DrawIndirect buffer during sort phase (GPU sort will override default DrawIndirect) diff --git a/tools/surfel_plane.rs b/tools/surfel_plane.rs new file mode 100644 index 00000000..a9b19fa6 --- /dev/null +++ b/tools/surfel_plane.rs @@ -0,0 +1,178 @@ +use bevy::{ + prelude::*, + app::AppExit, + core_pipeline::tonemapping::Tonemapping, +}; +use bevy_args::{ + BevyArgsPlugin, + parse_args, +}; +use bevy_inspector_egui::quick::WorldInspectorPlugin; +use bevy_panorbit_camera::{ + PanOrbitCamera, + PanOrbitCameraPlugin, +}; + +use bevy_gaussian_splatting::{ + Gaussian, + GaussianCamera, + GaussianCloud, + GaussianMode, + GaussianCloudSettings, + GaussianSplattingBundle, + GaussianSplattingPlugin, + utils::{ + setup_hooks, + GaussianSplattingViewer, + }, + SphericalHarmonicCoefficients, +}; + + +pub fn setup_surfel_compare( + mut commands: Commands, + mut gaussian_assets: ResMut>, +) { + let grid_size_x = 10; + let grid_size_y = 10; + let spacing = 5.0; + let visualize_bounding_box = false; + + let mut blue_gaussians = Vec::new(); + let mut blue_sh = SphericalHarmonicCoefficients::default(); + blue_sh.set(2, 5.0); + + for i in 0..grid_size_x { + for j in 0..grid_size_y { + let x = i as f32 * spacing - (grid_size_x as f32 * spacing) / 2.0; + let y = j as f32 * spacing - (grid_size_y as f32 * spacing) / 2.0; + let position = [x, y, 0.0, 1.0]; + let scale = [1.0, 1.0, 0.01, 0.5]; + + let gaussian = Gaussian { + position_visibility: position.into(), + rotation: [0.0, 0.0, 0.0, 1.0].into(), + scale_opacity: scale.into(), + spherical_harmonic: blue_sh, + }; + blue_gaussians.push(gaussian); + } + } + + commands.spawn(( + GaussianSplattingBundle { + cloud: gaussian_assets.add(GaussianCloud::from_gaussians(blue_gaussians)), + settings: GaussianCloudSettings { + visualize_bounding_box, + ..default() + }, + ..default() + }, + Name::new("gaussian_cloud_3dgs"), + )); + + let mut red_gaussians = Vec::new(); + let mut red_sh = SphericalHarmonicCoefficients::default(); + red_sh.set(0, 5.0); + + for i in 0..grid_size_x { + for j in 0..grid_size_y { + let x = i as f32 * spacing - (grid_size_x as f32 * spacing) / 2.0; + let y = j as f32 * spacing - (grid_size_y as f32 * spacing) / 2.0; + let position = [x, y, 0.0, 1.0]; + let scale = [1.0, 1.0, 99.0, 0.3]; + + // let angle = std::f32::consts::PI / 2.0; + // let rotation = Quat::from_rotation_y(angle).to_array().into(); + + let gaussian = Gaussian { + position_visibility: position.into(), + rotation: [0.0, 0.0, 0.0, 1.0].into(), + scale_opacity: scale.into(), + spherical_harmonic: red_sh, + }; + red_gaussians.push(gaussian); + } + } + + commands.spawn(( + GaussianSplattingBundle { + cloud: gaussian_assets.add(GaussianCloud::from_gaussians(red_gaussians)), + settings: GaussianCloudSettings { + visualize_bounding_box, + aabb: true, + transform: Transform::from_translation(Vec3::new(spacing, spacing, 0.0)), + gaussian_mode: GaussianMode::GaussianSurfel, + ..default() + }, + ..default() + }, + Name::new("gaussian_cloud_2dgs"), + )); + + commands.spawn(( + Camera3dBundle { + transform: Transform::from_translation(Vec3::new(0.0, 1.5, 20.0)), + tonemapping: Tonemapping::None, + ..default() + }, + PanOrbitCamera { + allow_upside_down: true, + ..default() + }, + GaussianCamera, + )); +} + +fn compare_surfel_app() { + let config = parse_args::(); + let mut app = App::new(); + + // setup for gaussian viewer app + app.insert_resource(ClearColor(Color::srgb_u8(0, 0, 0))); + app.add_plugins( + DefaultPlugins + .set(ImagePlugin::default_nearest()) + .set(WindowPlugin { + primary_window: Some(Window { + mode: bevy::window::WindowMode::Windowed, + present_mode: bevy::window::PresentMode::AutoVsync, + prevent_default_event_handling: false, + resolution: (config.width, config.height).into(), + title: config.name.clone(), + ..default() + }), + ..default() + }), + ); + app.add_plugins(BevyArgsPlugin::::default()); + app.add_plugins(PanOrbitCameraPlugin); + + if config.editor { + app.add_plugins(WorldInspectorPlugin::new()); + } + + if config.press_esc_close { + app.add_systems(Update, esc_close); + } + + // setup for gaussian splatting + app.add_plugins(GaussianSplattingPlugin); + app.add_systems(Startup, setup_surfel_compare); + + app.run(); +} + +pub fn esc_close( + keys: Res>, + mut exit: EventWriter +) { + if keys.just_pressed(KeyCode::Escape) { + exit.send(AppExit::Success); + } +} + +pub fn main() { + setup_hooks(); + compare_surfel_app(); +}