From f32c1ffdbdd2a4bab768f618ee01370e27630296 Mon Sep 17 00:00:00 2001 From: mosure Date: Wed, 22 Nov 2023 15:57:50 -0600 Subject: [PATCH 1/5] feat: aabb/obb compare, minor obb rotation error --- Cargo.toml | 5 ++ README.md | 1 + src/gaussian.rs | 2 + src/lib.rs | 1 + src/render/gaussian.wgsl | 93 ++++++++++----------- src/render/mod.rs | 2 +- tools/compare_aabb_obb.rs | 171 ++++++++++++++++++++++++++++++++++++++ viewer/viewer.rs | 12 ++- 8 files changed, 228 insertions(+), 59 deletions(-) create mode 100644 tools/compare_aabb_obb.rs diff --git a/Cargo.toml b/Cargo.toml index ee21d62d..a6633508 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -93,6 +93,11 @@ path = "tools/ply_to_gcloud.rs" required-features = ["io_ply"] +[[bin]] +name = "compare_aabb_obb" +path = "tools/compare_aabb_obb.rs" + + [[bench]] name = "io" harness = false diff --git a/README.md b/README.md index 87bece76..9d2a031d 100644 --- a/README.md +++ b/README.md @@ -63,6 +63,7 @@ fn setup_gaussian_cloud( - [ply to gcloud converter](tools/README.md#ply-to-gcloud-converter) - [ ] gaussian cloud training tool +- aabb vs. obb gaussian comparison via `cargo run --bin compare_aabb_obb` ## wasm support diff --git a/src/gaussian.rs b/src/gaussian.rs index 92682154..a65f4807 100644 --- a/src/gaussian.rs +++ b/src/gaussian.rs @@ -195,6 +195,8 @@ impl GaussianCloud { } } + cloud.0.push(cloud.0[0]); + cloud } } diff --git a/src/lib.rs b/src/lib.rs index 56ad130b..51bb1a66 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,6 +22,7 @@ pub mod utils; pub struct GaussianSplattingBundle { pub settings: GaussianCloudSettings, pub cloud: Handle, + pub visibility: Visibility, } diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index b71e1c97..e9f05a97 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -59,7 +59,11 @@ fn compute_cov3d(scale: vec3, rotation: vec4) -> array { ); } -fn compute_cov2d(position: vec3, scale: vec3, rotation: vec4) -> vec3 { +fn compute_cov2d( + position: vec3, + scale: vec3, + rotation: vec4 +) -> vec3 { let cov3d = compute_cov3d(scale, rotation); let Vrk = mat3x3( cov3d[0], cov3d[1], cov3d[2], @@ -69,30 +73,16 @@ fn compute_cov2d(position: vec3, scale: vec3, rotation: vec4) -> var t = view.inverse_view * vec4(position, 1.0); - let focal_x = 600.0; - let focal_y = 600.0; - - let fovy = 2.0 * atan(1.0 / view.projection[1][1]); - let fovx = 2.0 * atan(1.0 / view.projection[0][0]); - let tan_fovy = tan(fovy * 0.5); - let tan_fovx = tan(fovx * 0.5); - - let limx = 1.3 * tan_fovx; - let limy = 1.3 * tan_fovy; - let txtz = t.x / t.z; - let tytz = t.y / t.z; - t.x = min(limx, max(-limx, txtz)) * t.z; - t.y = min(limy, max(-limy, tytz)) * t.z; + let device_pixel_ratio = 1.0; + let focal = vec2( + view.projection.x.x * device_pixel_ratio * view.viewport.z * 0.45, + view.projection.y.y * device_pixel_ratio * view.viewport.w * 0.45, + ); + let s = 1.0 / (t.z * t.z); let J = mat3x3( - focal_x / t.z, - 0.0, - -(focal_x * t.x) / (t.z * t.z), - - 0.0, - -focal_y / t.z, - (focal_y * t.y) / (t.z * t.z), - + focal.x / t.z, 0.0, -(focal.x * t.x) * s, + 0.0, -focal.y / t.z, (focal.y * t.y) * s, 0.0, 0.0, 0.0, ); @@ -113,26 +103,22 @@ fn compute_cov2d(position: vec3, scale: vec3, rotation: vec4) -> return vec3(cov[0][0], cov[0][1], cov[1][1]); } - fn get_bounding_box( cov2d: vec3, direction: vec2, ) -> vec4 { // return vec4(offset, uv); - // let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y; - // let mid = 0.5 * (cov2d.x + cov2d.z); - // let lambda1 = mid + sqrt(max(0.1, mid * mid - det)); - // let lambda2 = mid - sqrt(max(0.1, mid * mid - det)); - // let x_axis_length = sqrt(lambda1); - // let y_axis_length = sqrt(lambda2); - let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y; - let mid = 0.5 * (cov2d.x + cov2d.z); - var discriminant = max(0.0, mid * mid - det); + let trace = cov2d.x + cov2d.z; + let mid = 0.5 * trace; + let discriminant = max(0.0, mid * mid - det); + + let term = sqrt(discriminant); + + let lambda1 = mid + term; + let lambda2 = max(mid - term, 0.0); - let lambda1 = mid + sqrt(discriminant); - let lambda2 = mid - sqrt(discriminant); let x_axis_length = sqrt(lambda1); let y_axis_length = sqrt(lambda2); @@ -146,7 +132,7 @@ fn get_bounding_box( ); return vec4( - 2.0 * radius_ndc * direction, + radius_ndc * direction, radius_px * direction, ); #endif @@ -161,31 +147,37 @@ fn get_bounding_box( // collapse unstable eigenvectors to circle let threshold = 0.1; if (abs(lambda1 - lambda2) < threshold) { - let circle = direction * max(x_axis_length, y_axis_length); + let circle = direction * max(bounds.x, bounds.y); return vec4( circle / view.viewport.zw, circle ); } - let eigvec1 = normalize(vec2( - cov2d.y, - lambda1 - cov2d.x + -cov2d.y, + lambda1 - cov2d.x, )); - let eigvec2 = vec2( - -eigvec1.y, - eigvec1.x - ); - - let rotation_matrix = mat2x2( - eigvec1.x, eigvec2.x, - eigvec1.y, eigvec2.y + let eigvec2 = normalize(vec2( + -cov2d.y, + lambda2 - cov2d.x, + )); + // let eigvec2 = vec2( + // eigvec1.y, + // -eigvec1.x + // ); + + let rotation_matrix = transpose( + mat2x2( + eigvec1, + eigvec2, + ) ); - let scaled_vertex = direction * bounds; + let scaling_factor = 1.0 / (0.5 * (view.viewport.z + view.viewport.w)); + let scaled_vertex = direction * bounds * scaling_factor; return vec4( - scaled_vertex * rotation_matrix / view.viewport.zw, + scaled_vertex * rotation_matrix, 0.0, 0.0, ); #endif @@ -234,7 +226,6 @@ fn vs_points( let cov2d = compute_cov2d(transformed_position, point.scale_opacity.rgb, point.rotation); - // TODO: remove conic when OBB is used let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y; let det_inv = 1.0 / det; let conic = vec3( diff --git a/src/render/mod.rs b/src/render/mod.rs index 2b93167a..ecde7d8a 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -831,7 +831,7 @@ pub fn queue_gaussian_bind_group( resource: BindingResource::Buffer(BufferBinding { buffer: model, offset: 0, - size: BufferSize::new(model.size()), + size: GaussianCloudUniform::min_size().into(), }), }, ], diff --git a/tools/compare_aabb_obb.rs b/tools/compare_aabb_obb.rs new file mode 100644 index 00000000..39fb0cee --- /dev/null +++ b/tools/compare_aabb_obb.rs @@ -0,0 +1,171 @@ +use bevy::{ + prelude::*, + app::AppExit, + core::Name, +}; +use bevy_inspector_egui::quick::WorldInspectorPlugin; +use bevy_panorbit_camera::{ + PanOrbitCamera, + PanOrbitCameraPlugin, +}; + +use bevy_gaussian_splatting::{ + Gaussian, + GaussianCloud, + GaussianCloudSettings, + GaussianSplattingBundle, + GaussianSplattingPlugin, + utils::setup_hooks, SphericalHarmonicCoefficients, +}; + + +// TODO: move to editor crate +pub struct GaussianSplattingViewer { + pub editor: bool, + pub esc_close: bool, + pub show_fps: bool, + pub width: f32, + pub height: f32, + pub name: String, +} + +impl Default for GaussianSplattingViewer { + fn default() -> GaussianSplattingViewer { + GaussianSplattingViewer { + editor: true, + esc_close: true, + show_fps: true, + width: 1920.0, + height: 1080.0, + name: "bevy_gaussian_splatting".to_string(), + } + } +} + + +pub fn setup_aabb_obb_compare( + mut commands: Commands, + mut gaussian_assets: ResMut>, +) { + let mut blue_sh = SphericalHarmonicCoefficients::default(); + blue_sh.coefficients[2] = 5.0; + + let blue_aabb_gaussian = Gaussian { + position: [0.0, 0.0, 0.0, 1.0], + rotation: [0.89, 0.0, -0.432, 0.144], + scale_opacity: [10.0, 1.0, 1.0, 0.5], + spherical_harmonic: blue_sh, + }; + + commands.spawn(( + GaussianSplattingBundle { + cloud: gaussian_assets.add( + GaussianCloud(vec![ + blue_aabb_gaussian, + blue_aabb_gaussian, + ]) + ), + settings: GaussianCloudSettings { + aabb: true, + visualize_bounding_box: true, + ..default() + }, + ..default() + }, + Name::new("gaussian_cloud_aabb"), + )); + + let mut red_sh = SphericalHarmonicCoefficients::default(); + red_sh.coefficients[0] = 5.0; + + let red_obb_gaussian = Gaussian { + position: [0.0, 0.0, 0.0, 1.0], + rotation: [0.89, 0.0, -0.432, 0.144], + scale_opacity: [10.0, 1.0, 1.0, 0.5], + spherical_harmonic: red_sh, + }; + + commands.spawn(( + GaussianSplattingBundle { + cloud: gaussian_assets.add( + GaussianCloud(vec![ + red_obb_gaussian, + red_obb_gaussian, + ]) + ), + settings: GaussianCloudSettings { + aabb: false, + visualize_bounding_box: true, + ..default() + }, + ..default() + }, + Name::new("gaussian_cloud_obb"), + )); + + commands.spawn(( + Camera3dBundle { + transform: Transform::from_translation(Vec3::new(0.0, 1.5, 5.0)), + ..default() + }, + PanOrbitCamera{ + allow_upside_down: true, + ..default() + }, + )); +} + +fn compare_aabb_obb_app() { + let config = GaussianSplattingViewer::default(); + let mut app = App::new(); + + // setup for gaussian viewer app + app.insert_resource(ClearColor(Color::rgb_u8(0, 0, 0))); + app.add_plugins( + DefaultPlugins + .set(ImagePlugin::default_nearest()) + .set(WindowPlugin { + primary_window: Some(Window { + fit_canvas_to_parent: false, + 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(( + PanOrbitCameraPlugin, + )); + + if config.editor { + app.add_plugins(WorldInspectorPlugin::new()); + } + + if config.esc_close { + app.add_systems(Update, esc_close); + } + + // setup for gaussian splatting + app.add_plugins(GaussianSplattingPlugin); + app.add_systems(Startup, setup_aabb_obb_compare); + + app.run(); +} + +pub fn esc_close( + keys: Res>, + mut exit: EventWriter +) { + if keys.just_pressed(KeyCode::Escape) { + exit.send(AppExit); + } +} + +pub fn main() { + setup_hooks(); + compare_aabb_obb_app(); +} diff --git a/viewer/viewer.rs b/viewer/viewer.rs index ed7ead62..a3ad39bc 100644 --- a/viewer/viewer.rs +++ b/viewer/viewer.rs @@ -15,7 +15,6 @@ use bevy_panorbit_camera::{ use bevy_gaussian_splatting::{ GaussianCloud, - GaussianCloudSettings, GaussianSplattingBundle, GaussianSplattingPlugin, random_gaussians, @@ -52,11 +51,6 @@ fn setup_gaussian_cloud( mut gaussian_assets: ResMut>, ) { let cloud: Handle; - let settings = GaussianCloudSettings { - aabb: true, - visualize_bounding_box: false, - ..default() - }; // TODO: add proper GaussianSplattingViewer argument parsing let file_arg = std::env::args().nth(1); @@ -64,6 +58,11 @@ fn setup_gaussian_cloud( println!("generating {} gaussians", n); cloud = gaussian_assets.add(random_gaussians(n)); } else if let Some(filename) = file_arg { + if filename == "--help".to_string() { + println!("usage: cargo run -- [filename | n]"); + return; + } + println!("loading {}", filename); cloud = asset_server.load(filename.to_string()); } else { @@ -73,7 +72,6 @@ fn setup_gaussian_cloud( commands.spawn(( GaussianSplattingBundle { cloud, - settings, ..default() }, Name::new("gaussian_cloud"), From aa3146f11f449b36eea501f08df0f1e9c857604a Mon Sep 17 00:00:00 2001 From: mosure Date: Wed, 22 Nov 2023 18:48:46 -0600 Subject: [PATCH 2/5] feat: parent visibility --- src/gaussian.rs | 4 +--- src/io/ply.rs | 3 ++- src/render/gaussian.wgsl | 18 +++++++----------- src/render/mod.rs | 15 ++++++++++++--- 4 files changed, 22 insertions(+), 18 deletions(-) diff --git a/src/gaussian.rs b/src/gaussian.rs index a65f4807..c2e76255 100644 --- a/src/gaussian.rs +++ b/src/gaussian.rs @@ -101,8 +101,6 @@ where } -pub const MAX_SIZE_VARIANCE: f32 = 5.0; - #[derive( Clone, Debug, @@ -215,7 +213,7 @@ impl Default for GaussianCloudSettings { fn default() -> Self { Self { aabb: false, - global_scale: 1.0, + global_scale: 2.0, global_transform: Transform::IDENTITY.into(), visualize_bounding_box: false, } diff --git a/src/io/ply.rs b/src/io/ply.rs index ab5d1da3..5db6ef1b 100644 --- a/src/io/ply.rs +++ b/src/io/ply.rs @@ -11,11 +11,12 @@ use ply_rs::{ use crate::gaussian::{ Gaussian, MAX_SH_COEFF_COUNT_PER_CHANNEL, - MAX_SIZE_VARIANCE, SH_CHANNELS, }; +pub const MAX_SIZE_VARIANCE: f32 = 5.0; + impl PropertyAccess for Gaussian { fn new() -> Self { Gaussian::default() diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index e9f05a97..19cde72c 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -150,7 +150,7 @@ fn get_bounding_box( let circle = direction * max(bounds.x, bounds.y); return vec4( circle / view.viewport.zw, - circle + circle, ); } @@ -158,14 +158,10 @@ fn get_bounding_box( -cov2d.y, lambda1 - cov2d.x, )); - let eigvec2 = normalize(vec2( - -cov2d.y, - lambda2 - cov2d.x, - )); - // let eigvec2 = vec2( - // eigvec1.y, - // -eigvec1.x - // ); + let eigvec2 = vec2( + eigvec1.y, + -eigvec1.x + ); let rotation_matrix = transpose( mat2x2( @@ -175,10 +171,10 @@ fn get_bounding_box( ); let scaling_factor = 1.0 / (0.5 * (view.viewport.z + view.viewport.w)); - let scaled_vertex = direction * bounds * scaling_factor; + let scaled_vertex = direction * bounds; return vec4( + scaled_vertex * rotation_matrix * scaling_factor, scaled_vertex * rotation_matrix, - 0.0, 0.0, ); #endif } diff --git a/src/render/mod.rs b/src/render/mod.rs index ecde7d8a..0546d535 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -766,6 +766,7 @@ pub fn extract_gaussians( Query<( Entity, // &ComputedVisibility, + &Visibility, &Handle, &GaussianCloudSettings, )>, @@ -774,7 +775,16 @@ pub fn extract_gaussians( let mut commands_list = Vec::with_capacity(*prev_commands_len); // let visible_gaussians = gaussians_query.iter().filter(|(_, vis, ..)| vis.is_visible()); - for (entity, verticies, settings) in gaussians_query.iter() { + for ( + entity, + visibility, + verticies, + settings, + ) in gaussians_query.iter() { + if visibility == Visibility::Hidden { + continue; + } + let settings_uniform = GaussianCloudUniform { transform: settings.global_transform.compute_matrix(), global_scale: settings.global_scale, @@ -805,7 +815,7 @@ pub struct GaussianCloudBindGroup { pub sorted_bind_group: BindGroup, } -pub fn queue_gaussian_bind_group( +fn queue_gaussian_bind_group( mut commands: Commands, mut groups: ResMut, gaussian_cloud_pipeline: Res, @@ -958,7 +968,6 @@ pub fn queue_gaussian_bind_group( } } - #[derive(Component)] pub struct GaussianViewBindGroup { pub value: BindGroup, From e12e2b6852eb3dc653f217884aff82cd5fe2ba7a Mon Sep 17 00:00:00 2001 From: mosure Date: Sat, 25 Nov 2023 17:05:27 -0600 Subject: [PATCH 3/5] default to AABB until OBB is resolved --- src/gaussian.rs | 2 +- src/render/gaussian.wgsl | 20 ++++++++++++-------- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/src/gaussian.rs b/src/gaussian.rs index c2e76255..0fb113e3 100644 --- a/src/gaussian.rs +++ b/src/gaussian.rs @@ -212,7 +212,7 @@ pub struct GaussianCloudSettings { impl Default for GaussianCloudSettings { fn default() -> Self { Self { - aabb: false, + aabb: true, global_scale: 2.0, global_transform: Transform::IDENTITY.into(), visualize_bounding_box: false, diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index 19cde72c..a2c39fbb 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -127,8 +127,7 @@ fn get_bounding_box( // creates a square AABB (inefficient fragment usage) let radius_px = 3.5 * max(x_axis_length, y_axis_length); let radius_ndc = vec2( - radius_px / view.viewport.z, - radius_px / view.viewport.w, + radius_px / view.viewport.zw, ); return vec4( @@ -138,12 +137,17 @@ fn get_bounding_box( #endif #ifdef USE_OBB + + let a = (cov2d.x - cov2d.z) * (cov2d.x - cov2d.z); + let b = sqrt(a + 4.0 * cov2d.y * cov2d.y); + 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( - x_axis_length, - y_axis_length, + major_radius, + minor_radius, ); - // bounding box is aligned to the eigenvectors with proper width/height // collapse unstable eigenvectors to circle let threshold = 0.1; if (abs(lambda1 - lambda2) < threshold) { @@ -170,10 +174,10 @@ fn get_bounding_box( ) ); - let scaling_factor = 1.0 / (0.5 * (view.viewport.z + view.viewport.w)); - let scaled_vertex = direction * bounds; + let scaling_factor = 1.0 / view.viewport.zw; + let scaled_vertex = direction * bounds * scaling_factor; return vec4( - scaled_vertex * rotation_matrix * scaling_factor, + scaled_vertex * rotation_matrix, scaled_vertex * rotation_matrix, ); #endif From 08b9d27a33aa0ca90759e43400e5ebc30b62b05a Mon Sep 17 00:00:00 2001 From: mosure Date: Sat, 25 Nov 2023 17:36:50 -0600 Subject: [PATCH 4/5] fix: OBB ndc calculation --- src/render/gaussian.wgsl | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index a2c39fbb..404276d4 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -174,11 +174,15 @@ fn get_bounding_box( ) ); + let scaled_vertex = direction * bounds; + let rotated_vertex = scaled_vertex * rotation_matrix; + let scaling_factor = 1.0 / view.viewport.zw; - let scaled_vertex = direction * bounds * scaling_factor; + let ndc_vertex = rotated_vertex * scaling_factor; + return vec4( - scaled_vertex * rotation_matrix, - scaled_vertex * rotation_matrix, + ndc_vertex, + rotated_vertex, ); #endif } @@ -252,8 +256,6 @@ fn vs_points( @fragment fn fs_main(input: GaussianOutput) -> @location(0) vec4 { - // TODO: draw gaussian without conic (OBB) - #ifdef USE_AABB let d = -input.major_minor; let conic = input.conic; From 828105ab6c3d49fe927124150ff957f6c2225152 Mon Sep 17 00:00:00 2001 From: mosure Date: Sat, 25 Nov 2023 17:37:13 -0600 Subject: [PATCH 5/5] default to OBB --- src/gaussian.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/gaussian.rs b/src/gaussian.rs index 0fb113e3..c2e76255 100644 --- a/src/gaussian.rs +++ b/src/gaussian.rs @@ -212,7 +212,7 @@ pub struct GaussianCloudSettings { impl Default for GaussianCloudSettings { fn default() -> Self { Self { - aabb: true, + aabb: false, global_scale: 2.0, global_transform: Transform::IDENTITY.into(), visualize_bounding_box: false,