feat: 4d gaussians (#32)

* refactor: separate compute nodes (resolves: #31) * checkin (broken morph pipeline) * merge from main * refactor: morph plugin * feat: extract particle behaviors * feat: particle behavior demo * docs: spatial queries as a future task (gaussian editor) * docs: clarify 4d behavior * fix: merge & disable cargo check * refactor: separate gaussian and sort pipelines * docs: more 4D related works * docs: add training pipeline tooling link * docs: cactus example * fix: include DrawIndirect
mosure · Nov 26, 2023 · 4356f87 · 4356f87
1 parent 5f626e9
commit 4356f87
Show file tree

Hide file tree

Showing 22 changed files with 1,323 additions and 592 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -33,9 +33,6 @@ jobs:
           target/
         key: ${{ runner.os }}-cargo-build-stable-${{ hashFiles('**/Cargo.toml') }}
 
-    - name: check
-      run: cargo check
-
     - name: build
       run: cargo build
 

diff --git a/Cargo.toml b/Cargo.toml
@@ -46,6 +46,15 @@ wasm-bindgen = "0.2"
 version = "0.12"
 default-features = true
 
+[dev-dependencies.bevy]
+version = "0.12"
+default-features = true
+features = [
+  'bevy_ci_testing',
+  'debug_glam_assert',
+]
+
+
 
 [dependencies.web-sys]
 version = "0.3.4"

diff --git a/README.md b/README.md
@@ -11,24 +11,28 @@
 bevy gaussian splatting render pipeline plugin
 
 ![Alt text](docs/notferris.png)
+![Alt text](docs/cactus.gif)
 ![Alt text](docs/bike.png)
 
-`cargo run -- scenes/icecream.gcloud`
+download [cactus.gcloud](https://mitchell.mosure.me/cactus.gcloud)
+
+`cargo run -- scenes/cactus.gcloud`
 
 ## capabilities
 
 - [X] ply to gcloud converter
 - [X] gcloud and ply asset loaders
 - [X] bevy gaussian cloud render pipeline
+- [X] 4D gaussian cloud wavelet compression
+- [X] gaussian cloud particle effects
+- [ ] accelerated spatial queries
 - [ ] wasm support /w [live demo](https://mosure.github.io/bevy_gaussian_splatting)
 - [ ] temporal depth sorting
 - [ ] f16 and f32 gcloud support
-- [ ] 4D gaussian clouds via morph targets
 - [ ] skeletons
 - [ ] volume masks
 - [ ] level of detail
 - [ ] lighting and shadows
-- [ ] gaussian cloud particle effects (accelerated spatial queries)
 - [ ] bevy_openxr support
 - [ ] bevy 3D camera to gaussian cloud pipeline
 
@@ -62,7 +66,7 @@ fn setup_gaussian_cloud(
 ## tools
 
 - [ply to gcloud converter](tools/README.md#ply-to-gcloud-converter)
-- [ ] gaussian cloud training tool
+- [gaussian cloud training pipeline](https://github.com/mosure/burn_gaussian_splatting)
 - aabb vs. obb gaussian comparison via `cargo run --bin compare_aabb_obb`
 
 ## wasm support
@@ -85,6 +89,7 @@ to build wasm run:
 - [4d gaussians](https://github.com/hustvl/4DGaussians)
 - [bevy](https://github.com/bevyengine/bevy)
 - [bevy-hanabi](https://github.com/djeedai/bevy_hanabi)
+- [d3ga](https://zielon.github.io/d3ga/)
 - [deformable-3d-gaussians](https://github.com/ingra14m/Deformable-3D-Gaussians)
 - [diff-gaussian-rasterization](https://github.com/graphdeco-inria/diff-gaussian-rasterization)
 - [dreamgaussian](https://github.com/dreamgaussian/dreamgaussian)
@@ -97,6 +102,7 @@ to build wasm run:
 - [masked-spacetime-hashing](https://github.com/masked-spacetime-hashing/msth)
 - [onesweep](https://arxiv.org/ftp/arxiv/papers/2206/2206.01784.pdf)
 - [pasture](https://github.com/Mortano/pasture)
+- [phys-gaussian](https://xpandora.github.io/PhysGaussian/)
 - [point-visualizer](https://github.com/mosure/point-visualizer)
 - [rusty-automata](https://github.com/mosure/rusty-automata)
 - [splat](https://github.com/antimatter15/splat)

diff --git a/assets/scenes/icecream.gcloud b/assets/scenes/icecream.gcloud
diff --git a/docs/cactus.gif b/docs/cactus.gif
diff --git a/src/gaussian.rs b/src/gaussian.rs
@@ -123,18 +123,22 @@ pub struct Gaussian {
     pub spherical_harmonic: SphericalHarmonicCoefficients,
 }
 
+
 #[derive(
     Asset,
     Clone,
     Debug,
+    Default,
     PartialEq,
     Reflect,
     TypeUuid,
     Serialize,
     Deserialize,
 )]
 #[uuid = "ac2f08eb-bc32-aabb-ff21-51571ea332d5"]
-pub struct GaussianCloud(pub Vec<Gaussian>);
+pub struct GaussianCloud {
+    pub gaussians: Vec<Gaussian>,
+}
 
 impl GaussianCloud {
     pub fn test_model() -> Self {
@@ -157,7 +161,7 @@ impl GaussianCloud {
                 0.5,
                 0.5,
             ],
-            spherical_harmonic: SphericalHarmonicCoefficients{
+            spherical_harmonic: SphericalHarmonicCoefficients {
                 coefficients: [
                     1.0, 0.0, 1.0,
                     0.0, 0.5, 0.0,
@@ -178,22 +182,25 @@ impl GaussianCloud {
                 ],
             },
         };
-        let mut cloud = GaussianCloud(Vec::new());
+        let mut cloud = GaussianCloud {
+            gaussians: Vec::new(),
+            ..default()
+        };
 
         for &x in [-0.5, 0.5].iter() {
             for &y in [-0.5, 0.5].iter() {
                 for &z in [-0.5, 0.5].iter() {
                     let mut g = origin;
                     g.position = [x, y, z, 1.0];
-                    cloud.0.push(g);
+                    cloud.gaussians.push(g);
 
                     let mut rng = rand::thread_rng();
-                    cloud.0.last_mut().unwrap().spherical_harmonic.coefficients.shuffle(&mut rng);
+                    cloud.gaussians.last_mut().unwrap().spherical_harmonic.coefficients.shuffle(&mut rng);
                 }
             }
         }
 
-        cloud.0.push(cloud.0[0]);
+        cloud.gaussians.push(cloud.gaussians[0]);
 
         cloud
     }
@@ -236,9 +243,9 @@ impl Distribution<Gaussian> for rand::distributions::Standard {
                 rng.gen_range(-1.0..1.0),
             ],
             scale_opacity: [
-                rng.gen_range(0.0..0.25),
-                rng.gen_range(0.0..0.25),
-                rng.gen_range(0.0..0.25),
+                rng.gen_range(0.0..1.0),
+                rng.gen_range(0.0..1.0),
+                rng.gen_range(0.0..1.0),
                 rng.gen_range(0.0..0.8),
             ],
             spherical_harmonic: SphericalHarmonicCoefficients {
@@ -260,5 +267,8 @@ pub fn random_gaussians(n: usize) -> GaussianCloud {
     for _ in 0..n {
         gaussians.push(rng.gen());
     }
-    GaussianCloud(gaussians)
+    GaussianCloud {
+        gaussians,
+        ..default()
+    }
 }
diff --git a/src/io/loader.rs b/src/io/loader.rs
@@ -46,8 +46,11 @@ impl AssetLoader for GaussianCloudLoader {
                         let cursor = Cursor::new(bytes);
                         let mut f = BufReader::new(cursor);
 
-                        let ply_cloud = crate::io::ply::parse_ply(&mut f)?;
-                        let cloud = GaussianCloud(ply_cloud);
+                        let gaussians = crate::io::ply::parse_ply(&mut f)?;
+                        let cloud = GaussianCloud {
+                            gaussians,
+                            ..Default::default()
+                        };
 
                         Ok(cloud)
                     }

diff --git a/src/render/bindings.wgsl b/src/render/bindings.wgsl
@@ -4,53 +4,32 @@
 #import bevy_render::view::View
 
 
-struct GaussianInput {
+@group(0) @binding(0) var<uniform> view: View;
+@group(0) @binding(1) var<uniform> globals: Globals;
+
+struct GaussianUniforms {
+    global_transform: mat4x4<f32>,
+    global_scale: f32,
+};
+@group(1) @binding(0) var<uniform> gaussian_uniforms: GaussianUniforms;
+
+struct Gaussian {
     @location(0) rotation: vec4<f32>,
     @location(1) position: vec4<f32>,
     @location(2) scale_opacity: vec4<f32>,
     sh: array<f32, #{MAX_SH_COEFF_COUNT}>,
 };
+@group(2) @binding(0) var<storage, read_write> points: array<Gaussian>;
 
-struct GaussianOutput {
-    @builtin(position) position: vec4<f32>,
-    @location(0) @interpolate(flat) color: vec4<f32>,
-    @location(1) @interpolate(flat) conic: vec3<f32>,
-    @location(2) @interpolate(linear) uv: vec2<f32>,
-    @location(3) @interpolate(linear) major_minor: vec2<f32>,
-};
-
-struct GaussianUniforms {
-    global_transform: mat4x4<f32>,
-    global_scale: f32,
-};
 
 struct DrawIndirect {
     vertex_count: u32,
     instance_count: atomic<u32>,
     base_vertex: u32,
     base_instance: u32,
 }
-struct SortingGlobal {
-    digit_histogram: array<array<atomic<u32>, #{RADIX_BASE}>, #{RADIX_DIGIT_PLACES}>,
-    assignment_counter: atomic<u32>,
-}
+
 struct Entry {
     key: u32,
     value: u32,
 }
-
-
-@group(0) @binding(0) var<uniform> view: View;
-@group(0) @binding(1) var<uniform> globals: Globals;
-
-@group(1) @binding(0) var<uniform> uniforms: GaussianUniforms;
-
-@group(2) @binding(0) var<storage, read> points: array<GaussianInput>;
-
-@group(3) @binding(0) var<uniform> sorting_pass_index: u32;
-@group(3) @binding(1) var<storage, read_write> sorting: SortingGlobal;
-@group(3) @binding(2) var<storage, read_write> status_counters: array<array<atomic<u32>, #{RADIX_BASE}>>;
-@group(3) @binding(3) var<storage, read_write> draw_indirect: DrawIndirect;
-@group(3) @binding(4) var<storage, read_write> input_entries: array<Entry>;
-@group(3) @binding(5) var<storage, read_write> output_entries: array<Entry>;
-@group(3) @binding(6) var<storage, read> sorted_entries: array<Entry>;
diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl
@@ -1,15 +1,14 @@
 #import bevy_gaussian_splatting::bindings::{
     view,
     globals,
-    uniforms,
+    gaussian_uniforms,
     points,
     sorting_pass_index,
     sorting,
     draw_indirect,
     input_entries,
     output_entries,
-    sorted_entries,
-    GaussianOutput,
+    Entry,
 }
 #import bevy_gaussian_splatting::spherical_harmonics::spherical_harmonics_lookup
 #import bevy_gaussian_splatting::transform::{
@@ -18,13 +17,24 @@
 }
 
 
+@group(3) @binding(0) var<storage, read> sorted_entries: array<Entry>;
+
+struct GaussianVertexOutput {
+    @builtin(position) position: vec4<f32>,
+    @location(0) @interpolate(flat) color: vec4<f32>,
+    @location(1) @interpolate(flat) conic: vec3<f32>,
+    @location(2) @interpolate(linear) uv: vec2<f32>,
+    @location(3) @interpolate(linear) major_minor: vec2<f32>,
+};
+
+
 // https://github.com/cvlab-epfl/gaussian-splatting-web/blob/905b3c0fb8961e42c79ef97e64609e82383ca1c2/src/shaders.ts#L185
 // TODO: precompute
 fn compute_cov3d(scale: vec3<f32>, rotation: vec4<f32>) -> array<f32, 6> {
     let S = mat3x3<f32>(
-        scale.x * uniforms.global_scale, 0.0, 0.0,
-        0.0, scale.y * uniforms.global_scale, 0.0,
-        0.0, 0.0, scale.z * uniforms.global_scale,
+        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,
     );
 
     let r = rotation.x;
@@ -192,8 +202,8 @@ fn get_bounding_box(
 fn vs_points(
     @builtin(instance_index) instance_index: u32,
     @builtin(vertex_index) vertex_index: u32,
-) -> GaussianOutput {
-    var output: GaussianOutput;
+) -> GaussianVertexOutput {
+    var output: GaussianVertexOutput;
     let splat_index = sorted_entries[instance_index][1];
 
     let discard_quad = sorted_entries[instance_index][0] == 0xFFFFFFFFu || splat_index == 0u;
@@ -204,7 +214,7 @@ fn vs_points(
     }
 
     let point = points[splat_index];
-    let transformed_position = (uniforms.global_transform * point.position).xyz;
+    let transformed_position = (gaussian_uniforms.global_transform * point.position).xyz;
     let projected_position = world_to_clip(transformed_position);
     if (!in_frustum(projected_position.xyz)) {
         output.color = vec4<f32>(0.0, 0.0, 0.0, 0.0);
@@ -228,6 +238,8 @@ fn vs_points(
         point.scale_opacity.a
     );
 
+    // TODO: add depth color visualization
+
     let cov2d = compute_cov2d(transformed_position, point.scale_opacity.rgb, point.rotation);
 
     let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y;
@@ -255,7 +267,7 @@ fn vs_points(
 }
 
 @fragment
-fn fs_main(input: GaussianOutput) -> @location(0) vec4<f32> {
+fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4<f32> {
 #ifdef USE_AABB
     let d = -input.major_minor;
     let conic = input.conic;