Skip to content

Commit

Permalink
feat: precompute cov3d (mosure#70)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
mosure authored Jan 9, 2024
1 parent 3b5013f commit eb0fcc3
Show file tree
Hide file tree
Showing 27 changed files with 631 additions and 108 deletions.
8 changes: 8 additions & 0 deletions .cargo/config.toml
Original file line number Diff line number Diff line change
Expand Up @@ -15,3 +15,11 @@ rustflags = [

[http]
proxy = ""


# offline development
# [source.crates-io]
# replace-with = "vendored-sources"

# [source.vendored-sources]
# directory = "vendor"
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
debug/
target/
vendor/
out/
Cargo.lock
**/*.rs.bk
Expand Down
14 changes: 11 additions & 3 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ exclude = [
default-run = "viewer"


# TODO: resolve one-hot feature flags through runtime configuration
[features]
default = [
"io_flexbuffers",
Expand All @@ -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",

Expand All @@ -72,6 +77,8 @@ noise = []
f32 = []
f16 = ["half"]

precompute_covariance_3d = [] # TODO: disable on main binary

packed = []
planar = []

Expand Down Expand Up @@ -101,6 +108,7 @@ web = [
"f16",
"io_flexbuffers",
"planar",
"precompute_covariance_3d",
"sort_std",
"viewer",
"webgl2",
Expand Down Expand Up @@ -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]
Expand Down
5 changes: 3 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
Binary file added docs/go.gif
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
83 changes: 75 additions & 8 deletions src/gaussian/cloud.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use rayon::prelude::*;
use crate::{
gaussian::{
f32::{
Covariance3dOpacity,
Position,
PositionVisibility,
Rotation,
Expand All @@ -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,
Expand All @@ -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>,
}

Expand Down Expand Up @@ -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];

Expand All @@ -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())
Expand Down Expand Up @@ -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());
}

Expand All @@ -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());
}
}
}
Expand All @@ -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,
}
}
Expand Down Expand Up @@ -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,
};

Expand Down
49 changes: 49 additions & 0 deletions src/gaussian/covariance.rs
Original file line number Diff line number Diff line change
@@ -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,
]
}
69 changes: 69 additions & 0 deletions src/gaussian/f16.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ use serde::{

use crate::gaussian::{
f32::{
Covariance3dOpacity,
Rotation,
ScaleOpacity,
},
Expand Down Expand Up @@ -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),
Expand Down
Loading

0 comments on commit eb0fcc3

Please sign in to comment.