feat: precompute cov3d (#70)

* feat: precompute_covariance_3d

* docs: credits

* docs: todo cleanup

* fix: clarify draw_indirect contents

* docs: use go example

* docs: remove bicycle from readme

* docs: credits

* chore: offline build support

* feat: precomputed covariance 3d - f16, buffer_storage

* fix: rotated covariance 3d

* fix: simplify frag uv

* docs: dynamic quad SH color

* docs: early depth termination

* docs: macos radix sort is unusable, MSAA may be too expensive

* fix: buffer_textures build

* docs: texture gcloud io

* fix: disable vendor mode

* fix: buffer_textures, precompute_cov3d support

.cargo/config.toml

@@ -15,3 +15,11 @@ rustflags = [
 
 [http]
 proxy = ""
+
+
+# offline development
+# [source.crates-io]
+# replace-with = "vendored-sources"
+
+# [source.vendored-sources]
+# directory = "vendor"

.gitignore

@@ -1,5 +1,6 @@
 debug/
 target/
+vendor/
 out/
 Cargo.lock
 **/*.rs.bk

Cargo.toml

@@ -32,6 +32,7 @@ exclude = [
 default-run = "viewer"
 
 
+# TODO: resolve one-hot feature flags through runtime configuration
 [features]
 default = [
   "io_flexbuffers",
@@ -46,10 +47,14 @@ default = [
   # "f32",
   "f16",
 
+  # "precompute_covariance_3d",
+
   "query_select",
-  "query_sparse",
+  # "query_sparse",
+
+  # "morph_particles",
 
-  "sort_radix",
+  # "sort_radix",  # TODO: fix macos radix sort
   "sort_rayon",
   "sort_std",
 
@@ -72,6 +77,8 @@ noise = []
 f32 = []
 f16 = ["half"]
 
+precompute_covariance_3d = []  # TODO: disable on main binary
+
 packed = []
 planar = []
 
@@ -101,6 +108,7 @@ web = [
   "f16",
   "io_flexbuffers",
   "planar",
+  "precompute_covariance_3d",
   "sort_std",
   "viewer",
   "webgl2",
@@ -134,7 +142,7 @@ wgpu = "0.17.1"
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 console_error_panic_hook = "0.1"
-wasm-bindgen = "0.2"
+wasm-bindgen = "0.2.89"
 
 
 [dependencies.bevy]

README.md

@@ -11,8 +11,7 @@
 bevy gaussian splatting render pipeline plugin. view the [live demo](https://mosure.github.io/bevy_gaussian_splatting/index.html?arg1=cactus.gcloud)
 
 ![Alt text](docs/notferris.png)
-![Alt text](docs/cactus.gif)
-![Alt text](docs/bike.png)
+![Alt text](docs/go.gif)
 
 
 ## capabilities
@@ -101,6 +100,8 @@ fn setup_gaussian_cloud(
 - [phys-gaussian](https://xpandora.github.io/PhysGaussian/)
 - [point-visualizer](https://github.com/mosure/point-visualizer)
 - [rusty-automata](https://github.com/mosure/rusty-automata)
+- [scaffold-gs](https://city-super.github.io/scaffold-gs/)
+- [shader-one-sweep](https://github.com/b0nes164/ShaderOneSweep)
 - [spacetime-gaussians](https://github.com/oppo-us-research/SpacetimeGaussians)
 - [splat](https://github.com/antimatter15/splat)
 - [splatter](https://github.com/Lichtso/splatter)

src/gaussian/cloud.rs

@@ -20,6 +20,7 @@ use rayon::prelude::*;
 use crate::{
     gaussian::{
         f32::{
+            Covariance3dOpacity,
             Position,
             PositionVisibility,
             Rotation,
@@ -34,13 +35,16 @@ use crate::{
     },
 };
 
+#[allow(unused_imports)]
 #[cfg(feature = "f16")]
 use crate::gaussian::f16::{
-        RotationScaleOpacityPacked128,
-        pack_f32s_to_u32,
-    };
+    Covariance3dOpacityPacked128,
+    RotationScaleOpacityPacked128,
+    pack_f32s_to_u32,
+};
 
 
+#[cfg(feature = "f16")]
 #[derive(
     Asset,
     Clone,
@@ -58,12 +62,37 @@ pub struct GaussianCloud {
 
     pub spherical_harmonic: Vec<SphericalHarmonicCoefficients>,
 
-    #[cfg(feature = "f16")]
+    #[cfg(not(feature = "precompute_covariance_3d"))]
     pub rotation_scale_opacity_packed128: Vec<RotationScaleOpacityPacked128>,
 
-    #[cfg(feature = "f32")]
+    #[cfg(feature = "precompute_covariance_3d")]
+    pub covariance_3d_opacity_packed128: Vec<Covariance3dOpacityPacked128>,
+}
+
+#[cfg(feature = "f32")]
+#[derive(
+    Asset,
+    Clone,
+    Debug,
+    Default,
+    PartialEq,
+    Reflect,
+    TypeUuid,
+    Serialize,
+    Deserialize,
+)]
+#[uuid = "ac2f08eb-bc32-aabb-ff21-51571ea332d5"]
+pub struct GaussianCloud {
+    pub position_visibility: Vec<PositionVisibility>,
+
+    pub spherical_harmonic: Vec<SphericalHarmonicCoefficients>,
+
+    #[cfg(feature = "precompute_covariance_3d")]
+    pub covariance_3d: Vec<Covariance3dOpacity>,
+
+    #[cfg(not(feature = "precompute_covariance_3d"))]
     pub rotation: Vec<Rotation>,
-    #[cfg(feature = "f32")]
+    #[cfg(not(feature = "precompute_covariance_3d"))]
     pub scale_opacity: Vec<ScaleOpacity>,
 }
 
@@ -146,7 +175,10 @@ impl GaussianCloud {
     //     return &mut self.scale_opacity[index].scale;
     // }
 
-    #[cfg(feature = "f16")]
+    #[cfg(all(
+        not(feature = "precompute_covariance_3d"),
+        feature = "f16",
+    ))]
     pub fn gaussian(&self, index: usize) -> Gaussian {
         let rso = self.rotation_scale_opacity_packed128[index];
 
@@ -171,7 +203,10 @@ impl GaussianCloud {
         }
     }
 
-    #[cfg(feature = "f16")]
+    #[cfg(all(
+        not(feature = "precompute_covariance_3d"),
+        feature = "f16",
+    ))]
     pub fn gaussian_iter(&self) -> impl Iterator<Item=Gaussian> + '_ {
         self.position_visibility.iter()
             .zip(self.spherical_harmonic.iter())
@@ -218,6 +253,10 @@ impl GaussianCloud {
         {
             self.position_visibility.resize(self.square_len(), PositionVisibility::default());
             self.spherical_harmonic.resize(self.square_len(), SphericalHarmonicCoefficients::default());
+
+            #[cfg(feature = "precompute_covariance_3d")]
+            self.covariance_3d_opacity_packed128.resize(self.square_len(), Covariance3dOpacityPacked128::default());
+            #[cfg(not(feature = "precompute_covariance_3d"))]
             self.rotation_scale_opacity_packed128.resize(self.square_len(), RotationScaleOpacityPacked128::default());
         }
 
@@ -227,6 +266,7 @@ impl GaussianCloud {
             self.spherical_harmonic.resize(self.square_len(), SphericalHarmonicCoefficients::default());
             self.rotation.resize(self.square_len(), Rotation::default());
             self.scale_opacity.resize(self.square_len(), ScaleOpacity::default());
+            self.covariance_3d.resize(self.square_len(), Covariance3dOpacity::default());
         }
     }
 }
@@ -237,17 +277,31 @@ impl GaussianCloud {
     pub fn subset(&self, indicies: &[usize]) -> Self {
         let mut position_visibility = Vec::with_capacity(indicies.len());
         let mut spherical_harmonic = Vec::with_capacity(indicies.len());
+
+        #[cfg(feature = "precompute_covariance_3d")]
+        let mut covariance_3d_opacity_packed128 = Vec::with_capacity(indicies.len());
+
+        #[cfg(not(feature = "precompute_covariance_3d"))]
         let mut rotation_scale_opacity_packed128 = Vec::with_capacity(indicies.len());
 
         for &index in indicies.iter() {
             position_visibility.push(self.position_visibility[index]);
             spherical_harmonic.push(self.spherical_harmonic[index]);
+
+            #[cfg(feature = "precompute_covariance_3d")]
+            covariance_3d_opacity_packed128.push(self.covariance_3d_opacity_packed128[index]);
+
+            #[cfg(not(feature = "precompute_covariance_3d"))]
             rotation_scale_opacity_packed128.push(self.rotation_scale_opacity_packed128[index]);
         }
 
         Self {
             position_visibility,
             spherical_harmonic,
+
+            #[cfg(feature = "precompute_covariance_3d")]
+            covariance_3d_opacity_packed128,
+            #[cfg(not(feature = "precompute_covariance_3d"))]
             rotation_scale_opacity_packed128,
         }
     }
@@ -292,19 +346,32 @@ impl GaussianCloud {
     pub fn from_gaussians(gaussians: Vec<Gaussian>) -> Self {
         let mut position_visibility = Vec::with_capacity(gaussians.len());
         let mut spherical_harmonic = Vec::with_capacity(gaussians.len());
+
+        #[cfg(feature = "precompute_covariance_3d")]
+        let mut covariance_3d_opacity_packed128 = Vec::with_capacity(gaussians.len());
+
+        #[cfg(not(feature = "precompute_covariance_3d"))]
         let mut rotation_scale_opacity_packed128 = Vec::with_capacity(gaussians.len());
 
         for gaussian in gaussians {
             position_visibility.push(gaussian.position_visibility);
             spherical_harmonic.push(gaussian.spherical_harmonic);
 
+            #[cfg(feature = "precompute_covariance_3d")]
+            covariance_3d_opacity_packed128.push(Covariance3dOpacityPacked128::from_gaussian(&gaussian));
+
+            #[cfg(not(feature = "precompute_covariance_3d"))]
             rotation_scale_opacity_packed128.push(RotationScaleOpacityPacked128::from_gaussian(&gaussian));
         }
 
         #[allow(unused_mut)]
         let mut cloud = GaussianCloud {
             position_visibility,
             spherical_harmonic,
+
+            #[cfg(feature = "precompute_covariance_3d")]
+            covariance_3d_opacity_packed128,
+            #[cfg(not(feature = "precompute_covariance_3d"))]
             rotation_scale_opacity_packed128,
         };
 

src/gaussian/covariance.rs

@@ -0,0 +1,49 @@
+use bevy::math::{
+    Mat3,
+    Vec3,
+    Vec4,
+};
+
+
+#[allow(non_snake_case)]
+pub fn compute_covariance_3d(
+    rotation: Vec4,
+    scale: Vec3,
+) -> [f32; 6] {
+    let S = Mat3::from_diagonal(scale);
+
+    let r = rotation.x;
+    let x = rotation.y;
+    let y = rotation.z;
+    let z = rotation.w;
+
+    let R = Mat3::from_cols(
+        Vec3::new(
+            1.0 - 2.0 * (y * y + z * z),
+            2.0 * (x * y - r * z),
+            2.0 * (x * z + r * y),
+        ),
+        Vec3::new(
+            2.0 * (x * y + r * z),
+            1.0 - 2.0 * (x * x + z * z),
+            2.0 * (y * z - r * x),
+        ),
+        Vec3::new(
+            2.0 * (x * z - r * y),
+            2.0 * (y * z + r * x),
+            1.0 - 2.0 * (x * x + y * y),
+        ),
+    );
+
+    let M = S * R;
+    let Sigma = M.transpose() * M;
+
+    [
+        Sigma.row(0).x,
+        Sigma.row(0).y,
+        Sigma.row(0).z,
+        Sigma.row(1).y,
+        Sigma.row(1).z,
+        Sigma.row(2).z,
+    ]
+}

src/gaussian/f16.rs

@@ -17,6 +17,7 @@ use serde::{
 
 use crate::gaussian::{
     f32::{
+        Covariance3dOpacity,
         Rotation,
         ScaleOpacity,
     },
@@ -128,6 +129,74 @@ impl From<[u32; 4]> for RotationScaleOpacityPacked128 {
 }
 
 
+#[derive(
+    Clone,
+    Debug,
+    Default,
+    Copy,
+    PartialEq,
+    Reflect,
+    ShaderType,
+    Pod,
+    Zeroable,
+    Serialize,
+    Deserialize,
+)]
+#[repr(C)]
+pub struct Covariance3dOpacityPacked128 {
+    #[reflect(ignore)]
+    pub cov3d: [u32; 3],
+    pub opacity: u32,
+}
+
+impl Covariance3dOpacityPacked128 {
+    pub fn from_gaussian(gaussian: &Gaussian) -> Self {
+        let cov3d: Covariance3dOpacity = gaussian.into();
+        let cov3d = cov3d.cov3d;
+
+        let opacity = gaussian.scale_opacity.opacity;
+
+        Self {
+            cov3d: [
+                pack_f32s_to_u32(cov3d[0], cov3d[1]),
+                pack_f32s_to_u32(cov3d[2], cov3d[3]),
+                pack_f32s_to_u32(cov3d[4], cov3d[5]),
+            ],
+            opacity: pack_f32s_to_u32(opacity, opacity),  // TODO: benefit from 32-bit opacity
+        }
+    }
+
+    pub fn covariance_3d_opacity(&self) -> Covariance3dOpacity {
+        let (c0, c1) = unpack_u32_to_f32s(self.cov3d[0]);
+        let (c2, c3) = unpack_u32_to_f32s(self.cov3d[1]);
+        let (c4, c5) = unpack_u32_to_f32s(self.cov3d[2]);
+
+        let (opacity, _) = unpack_u32_to_f32s(self.opacity);
+
+        let cov3d: [f32; 6] = [c0, c1, c2, c3, c4, c5];
+
+        Covariance3dOpacity {
+            cov3d,
+            opacity,
+            pad: 0.0,
+        }
+    }
+}
+
+impl From<[u32; 4]> for Covariance3dOpacityPacked128 {
+    fn from(cov3d_opacity: [u32; 4]) -> Self {
+        Self {
+            cov3d: [
+                cov3d_opacity[0],
+                cov3d_opacity[1],
+                cov3d_opacity[2],
+            ],
+            opacity: cov3d_opacity[3],
+        }
+    }
+}
+
+
 pub fn pack_f32s_to_u32(upper: f32, lower: f32) -> u32 {
     pack_f16s_to_u32(
         f16::from_f32(upper),