feat: aabb/obb compare (#36)

* feat: aabb/obb compare, minor obb rotation error

* feat: parent visibility

* default to AABB until OBB is resolved

* fix: OBB ndc calculation

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

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
 

src/gaussian.rs

@@ -101,8 +101,6 @@ where
 }
 
 
-pub const MAX_SIZE_VARIANCE: f32 = 5.0;
-
 #[derive(
     Clone,
     Debug,
@@ -195,6 +193,8 @@ impl GaussianCloud {
             }
         }
 
+        cloud.0.push(cloud.0[0]);
+
         cloud
     }
 }
@@ -213,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,
         }

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()

src/lib.rs

@@ -22,6 +22,7 @@ pub mod utils;
 pub struct GaussianSplattingBundle {
     pub settings: GaussianCloudSettings,
     pub cloud: Handle<GaussianCloud>,
+    pub visibility: Visibility,
 }
 
 

src/render/gaussian.wgsl

@@ -59,7 +59,11 @@ fn compute_cov3d(scale: vec3<f32>, rotation: vec4<f32>) -> array<f32, 6> {
     );
 }
 
-fn compute_cov2d(position: vec3<f32>, scale: vec3<f32>, rotation: vec4<f32>) -> vec3<f32> {
+fn compute_cov2d(
+    position: vec3<f32>,
+    scale: vec3<f32>,
+    rotation: vec4<f32>
+) -> vec3<f32> {
     let cov3d = compute_cov3d(scale, rotation);
     let Vrk = mat3x3(
         cov3d[0], cov3d[1], cov3d[2],
@@ -69,30 +73,16 @@ fn compute_cov2d(position: vec3<f32>, scale: vec3<f32>, rotation: vec4<f32>) ->
 
     var t = view.inverse_view * vec4<f32>(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<f32>(
+        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<f32>, scale: vec3<f32>, rotation: vec4<f32>) ->
     return vec3<f32>(cov[0][0], cov[0][1], cov[1][1]);
 }
 
-
 fn get_bounding_box(
     cov2d: vec3<f32>,
     direction: vec2<f32>,
 ) -> vec4<f32> {
     // return vec4<f32>(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);
 
@@ -141,52 +127,62 @@ 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<f32>(
-        radius_px / view.viewport.z,
-        radius_px / view.viewport.w,
+        radius_px / view.viewport.zw,
     );
 
     return vec4<f32>(
-        2.0 * radius_ndc * direction,
+        radius_ndc * direction,
         radius_px * direction,
     );
 #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<f32>(
-        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) {
-        let circle = direction * max(x_axis_length, y_axis_length);
+        let circle = direction * max(bounds.x, bounds.y);
         return vec4<f32>(
             circle / view.viewport.zw,
-            circle
+            circle,
         );
     }
 
-
     let eigvec1 = normalize(vec2<f32>(
-        cov2d.y,
-        lambda1 - cov2d.x
+        -cov2d.y,
+        lambda1 - cov2d.x,
     ));
     let eigvec2 = vec2<f32>(
-        -eigvec1.y,
-        eigvec1.x
+        eigvec1.y,
+        -eigvec1.x
     );
 
-    let rotation_matrix = mat2x2(
-        eigvec1.x, eigvec2.x,
-        eigvec1.y, eigvec2.y
+    let rotation_matrix = transpose(
+        mat2x2(
+            eigvec1,
+            eigvec2,
+        )
     );
 
     let scaled_vertex = direction * bounds;
+    let rotated_vertex = scaled_vertex * rotation_matrix;
+
+    let scaling_factor = 1.0 / view.viewport.zw;
+    let ndc_vertex = rotated_vertex * scaling_factor;
+
     return vec4<f32>(
-        scaled_vertex * rotation_matrix / view.viewport.zw,
-        0.0, 0.0,
+        ndc_vertex,
+        rotated_vertex,
     );
 #endif
 }
@@ -234,7 +230,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<f32>(
@@ -261,8 +256,6 @@ fn vs_points(
 
 @fragment
 fn fs_main(input: GaussianOutput) -> @location(0) vec4<f32> {
-    // TODO: draw gaussian without conic (OBB)
-
 #ifdef USE_AABB
     let d = -input.major_minor;
     let conic = input.conic;

src/render/mod.rs

@@ -766,6 +766,7 @@ pub fn extract_gaussians(
         Query<(
             Entity,
             // &ComputedVisibility,
+            &Visibility,
             &Handle<GaussianCloud>,
             &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<GaussianUniformBindGroups>,
     gaussian_cloud_pipeline: Res<GaussianCloudPipeline>,
@@ -831,7 +841,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(),
                 }),
             },
         ],
@@ -958,7 +968,6 @@ pub fn queue_gaussian_bind_group(
     }
 }
 
-
 #[derive(Component)]
 pub struct GaussianViewBindGroup {
     pub value: BindGroup,