diff --git a/.github/workflows/todo_tracker.yml b/.github/workflows/todo_tracker.yml index c2f9eedb..a3b1eb6d 100644 --- a/.github/workflows/todo_tracker.yml +++ b/.github/workflows/todo_tracker.yml @@ -7,6 +7,9 @@ on: jobs: build: + permissions: + issues: write + name: todo_tracker runs-on: [ubuntu-latest] diff --git a/.gitignore b/.gitignore index 7012bb23..a64aac48 100644 --- a/.gitignore +++ b/.gitignore @@ -6,5 +6,6 @@ Cargo.lock *.ply +*.gcloud .DS_Store diff --git a/Cargo.toml b/Cargo.toml index 64487ef2..3f968c83 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,14 +10,20 @@ categories = ["computer-vision", "graphics", "rendering", "rendering::data-forma homepage = "https://github.com/mosure/bevy_gaussian_splatting" repository = "https://github.com/mosure/bevy_gaussian_splatting" readme = "README.md" +include = ["tools"] exclude = [".devcontainer", ".github", "docs", "dist", "build", "assets", "credits"] +default-run = "bevy_gaussian_splatting" # TODO: use minimal bevy features [dependencies] -bevy = "0.11.2" +bevy = "0.11.3" +bevy-inspector-egui = "0.20.0" bevy_panorbit_camera = "0.8.0" +bincode2 = "2.0.1" bytemuck = "1.14.0" +flate2 = "1.0.28" ply-rs = "0.1.3" +serde = "1.0.189" [target.'cfg(target_arch = "wasm32")'.dependencies] @@ -51,3 +57,15 @@ inherits = "release" opt-level = "z" lto = "fat" codegen-units = 1 + + +[lib] +path = "src/lib.rs" + +[[bin]] +name = "bevy_gaussian_splatting" +path = "src/main.rs" + +[[bin]] +name = "ply_to_gcloud" +path = "tools/ply_to_gcloud.rs" diff --git a/README.md b/README.md index d9f7905b..05e56175 100644 --- a/README.md +++ b/README.md @@ -12,9 +12,13 @@ bevy gaussian splatting render pipeline plugin +`cargo run -- {path to ply or gcloud.gz file}` + ## capabilities -- [ ] bevy gaussian cloud render pipeline +- [X] ply to gcloud converter +- [X] gcloud and ply asset loaders +- [X] bevy gaussian cloud render pipeline - [ ] 4D gaussian clouds via morph targets - [ ] bevy 3D camera to gaussian cloud pipeline @@ -55,11 +59,13 @@ fn setup_gaussian_cloud( # credits +- [4d gaussians](https://github.com/hustvl/4DGaussians) - [bevy](https://github.com/bevyengine/bevy) - [bevy-hanabi](https://github.com/djeedai/bevy_hanabi) - [diff-gaussian-rasterization](https://github.com/graphdeco-inria/diff-gaussian-rasterization) - [dreamgaussian](https://github.com/dreamgaussian/dreamgaussian) - [dynamic-3d-gaussians](https://github.com/JonathonLuiten/Dynamic3DGaussians) +- [ewa splatting](https://www.cs.umd.edu/~zwicker/publications/EWASplatting-TVCG02.pdf) - [gaussian-splatting](https://github.com/graphdeco-inria/gaussian-splatting) - [gaussian-splatting-web](https://github.com/cvlab-epfl/gaussian-splatting-web) - [making gaussian splats smaller](https://aras-p.info/blog/2023/09/13/Making-Gaussian-Splats-smaller/) diff --git a/assets/scenes/icecream.gcloud b/assets/scenes/icecream.gcloud new file mode 100644 index 00000000..6828e088 Binary files /dev/null and b/assets/scenes/icecream.gcloud differ diff --git a/src/gaussian.rs b/src/gaussian.rs index d9742220..00effede 100644 --- a/src/gaussian.rs +++ b/src/gaussian.rs @@ -13,17 +13,22 @@ use bevy::{ LoadContext, LoadedAsset, }, - reflect::{ - TypePath, - TypeUuid, - }, + reflect::TypeUuid, render::render_resource::ShaderType, utils::BoxedFuture, }; +use bincode2::deserialize_from; use bytemuck::{ Pod, Zeroable, }; +use flate2::read::GzDecoder; +use serde::{ + Deserialize, + Serialize, + Serializer, + ser::SerializeTuple, +}; use crate::ply::parse_ply; @@ -37,9 +42,19 @@ const fn num_sh_coefficients(degree: usize) -> usize { } const SH_DEGREE: usize = 3; pub const MAX_SH_COEFF_COUNT: usize = num_sh_coefficients(SH_DEGREE) * 3; -#[derive(Clone, Copy, ShaderType, Pod, Zeroable)] +#[derive( + Clone, + Copy, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] #[repr(C)] pub struct SphericalHarmonicCoefficients { + #[serde(serialize_with = "coefficients_serializer", deserialize_with = "coefficients_deserializer")] pub coefficients: [f32; MAX_SH_COEFF_COUNT], } impl Default for SphericalHarmonicCoefficients { @@ -49,12 +64,64 @@ impl Default for SphericalHarmonicCoefficients { } } } +fn coefficients_serializer(n: &[f32; MAX_SH_COEFF_COUNT], s: S) -> Result +where + S: Serializer, +{ + let mut tup = s.serialize_tuple(MAX_SH_COEFF_COUNT)?; + for &x in n.iter() { + tup.serialize_element(&x)?; + } + + tup.end() +} + +fn coefficients_deserializer<'de, D>(d: D) -> Result<[f32; MAX_SH_COEFF_COUNT], D::Error> +where + D: serde::Deserializer<'de>, +{ + struct CoefficientsVisitor; + + impl<'de> serde::de::Visitor<'de> for CoefficientsVisitor { + type Value = [f32; MAX_SH_COEFF_COUNT]; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("an array of floats") + } + + fn visit_seq(self, mut seq: A) -> Result<[f32; MAX_SH_COEFF_COUNT], A::Error> + where + A: serde::de::SeqAccess<'de>, + { + let mut coefficients = [0.0; MAX_SH_COEFF_COUNT]; + for i in 0..MAX_SH_COEFF_COUNT { + coefficients[i] = seq + .next_element()? + .ok_or_else(|| serde::de::Error::invalid_length(i, &self))?; + } + Ok(coefficients) + } + } + + d.deserialize_tuple(MAX_SH_COEFF_COUNT, CoefficientsVisitor) +} + -#[derive(Clone, Default, Copy, ShaderType, Pod, Zeroable)] +pub const MAX_SIZE_VARIANCE: f32 = 5.0; + +#[derive( + Clone, + Default, + Copy, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] #[repr(C)] pub struct Gaussian { - //pub anisotropic_covariance: AnisotropicCovariance, - //pub normal: Vec3, pub rotation: [f32; 4], pub position: Vec3, pub scale: Vec3, @@ -63,7 +130,13 @@ pub struct Gaussian { padding: f32, } -#[derive(Clone, TypeUuid, TypePath)] +#[derive( + Clone, + Reflect, + TypeUuid, + Serialize, + Deserialize, +)] #[uuid = "ac2f08eb-bc32-aabb-ff21-51571ea332d5"] pub struct GaussianCloud(pub Vec); @@ -103,9 +176,9 @@ impl GaussianCloud { }; let mut cloud = GaussianCloud(Vec::new()); - for &x in [-1.0, 1.0].iter() { - for &y in [-1.0, 1.0].iter() { - for &z in [-1.0, 1.0].iter() { + 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.clone(); g.position = Vec3::new(x, y, z); cloud.0.push(g); @@ -120,15 +193,19 @@ impl GaussianCloud { #[derive(Component, Reflect, Clone)] pub struct GaussianCloudSettings { + pub aabb: bool, pub global_scale: f32, pub global_transform: GlobalTransform, + pub visualize_bounding_box: bool, } impl Default for GaussianCloudSettings { fn default() -> Self { Self { + aabb: false, global_scale: 1.0, global_transform: Transform::IDENTITY.into(), + visualize_bounding_box: false, } } } @@ -144,18 +221,30 @@ impl AssetLoader for GaussianCloudLoader { load_context: &'a mut LoadContext, ) -> BoxedFuture<'a, Result<(), bevy::asset::Error>> { Box::pin(async move { - let cursor = Cursor::new(bytes); - let mut f = BufReader::new(cursor); - - let ply_cloud = parse_ply(&mut f)?; - let cloud = GaussianCloud(ply_cloud); - - load_context.set_default_asset(LoadedAsset::new(cloud)); - Ok(()) + match load_context.path().extension() { + Some(ext) if ext == "ply" => { + let cursor = Cursor::new(bytes); + let mut f = BufReader::new(cursor); + + let ply_cloud = parse_ply(&mut f)?; + let cloud = GaussianCloud(ply_cloud); + + load_context.set_default_asset(LoadedAsset::new(cloud)); + return Ok(()); + }, + Some(ext) if ext == "gcloud" => { + let decompressed = GzDecoder::new(bytes); + let cloud: GaussianCloud = deserialize_from(decompressed).expect("failed to decode cloud"); + + load_context.set_default_asset(LoadedAsset::new(cloud)); + return Ok(()); + }, + _ => Ok(()), + } }) } fn extensions(&self) -> &[&str] { - &["ply"] + &["ply", "gcloud"] } } diff --git a/src/lib.rs b/src/lib.rs index e6a33107..fd542201 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,7 +18,7 @@ pub mod utils; #[derive(Bundle, Default, Reflect)] pub struct GaussianSplattingBundle { - pub settings: GaussianCloudSettings, // TODO: implement global transform + pub settings: GaussianCloudSettings, pub cloud: Handle, } @@ -31,9 +31,12 @@ pub struct GaussianSplattingPlugin; impl Plugin for GaussianSplattingPlugin { fn build(&self, app: &mut App) { + // TODO: allow hot reloading of GaussianCloud handle through inspector UI app.add_asset::(); app.init_asset_loader::(); + app.register_asset_reflect::(); + app.register_type::(); app.register_type::(); app.add_plugins(( diff --git a/src/main.rs b/src/main.rs index 3f8cc1f0..9d17f319 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,6 +6,7 @@ use bevy::{ FrameTimeDiagnosticsPlugin, }, }; +use bevy_inspector_egui::quick::WorldInspectorPlugin; use bevy_panorbit_camera::{ PanOrbitCamera, PanOrbitCameraPlugin, @@ -13,6 +14,7 @@ use bevy_panorbit_camera::{ use bevy_gaussian_splatting::{ GaussianCloud, + GaussianCloudSettings, GaussianSplattingBundle, GaussianSplattingPlugin, utils::setup_hooks, @@ -20,6 +22,7 @@ use bevy_gaussian_splatting::{ pub struct GaussianSplattingViewer { + pub editor: bool, pub esc_close: bool, pub show_fps: bool, pub width: f32, @@ -30,6 +33,7 @@ pub struct GaussianSplattingViewer { impl Default for GaussianSplattingViewer { fn default() -> GaussianSplattingViewer { GaussianSplattingViewer { + editor: true, esc_close: true, show_fps: true, width: 1920.0, @@ -42,14 +46,27 @@ impl Default for GaussianSplattingViewer { fn setup_gaussian_cloud( mut commands: Commands, - _asset_server: Res, + asset_server: Res, mut gaussian_assets: ResMut>, ) { - let cloud = gaussian_assets.add(GaussianCloud::test_model()); + let cloud: Handle; + let settings = GaussianCloudSettings { + aabb: true, + visualize_bounding_box: false, + ..default() + }; + + let filename = std::env::args().nth(1); + if let Some(filename) = filename { + println!("loading {}", filename); + cloud = asset_server.load(filename.as_str()); + } else { + cloud = gaussian_assets.add(GaussianCloud::test_model()); + } + commands.spawn(GaussianSplattingBundle { cloud, - // cloud: _asset_server.load("scenes/icecream.ply"), - ..Default::default() + settings, }); commands.spawn(( @@ -82,12 +99,16 @@ fn example_app() { ..default() }), ..default() - }) + }), ); app.add_plugins(( PanOrbitCameraPlugin, )); + if config.editor { + app.add_plugins(WorldInspectorPlugin::new()); + } + if config.esc_close { app.add_systems(Update, esc_close); } diff --git a/src/ply.rs b/src/ply.rs index ab9fa032..2a64d61a 100644 --- a/src/ply.rs +++ b/src/ply.rs @@ -1,6 +1,9 @@ use std::io::BufRead; -use bevy::asset::Error; +use bevy::{ + asset::Error, + math::Vec3, +}; use ply_rs::{ ply::{ Property, @@ -9,7 +12,10 @@ use ply_rs::{ parser::Parser, }; -use crate::gaussian::Gaussian; +use crate::gaussian::{ + Gaussian, + MAX_SIZE_VARIANCE, +}; impl PropertyAccess for Gaussian { @@ -29,9 +35,9 @@ impl PropertyAccess for Gaussian { ("f_dc_1", Property::Float(v)) => self.spherical_harmonic.coefficients[1] = v, ("f_dc_2", Property::Float(v)) => self.spherical_harmonic.coefficients[2] = v, ("opacity", Property::Float(v)) => self.opacity = 1.0 / (1.0 + (-v).exp()), - ("scale_0", Property::Float(v)) => self.scale.x = v.exp(), // TODO: variance cap: https://github.com/Lichtso/splatter/blob/c6b7a3894c25578cd29c9761619e4f194449e389/src/scene.rs#L235 - ("scale_1", Property::Float(v)) => self.scale.y = v.exp(), - ("scale_2", Property::Float(v)) => self.scale.z = v.exp(), + ("scale_0", Property::Float(v)) => self.scale.x = v, + ("scale_1", Property::Float(v)) => self.scale.y = v, + ("scale_2", Property::Float(v)) => self.scale.z = v, ("rot_0", Property::Float(v)) => self.rotation[0] = v, ("rot_1", Property::Float(v)) => self.rotation[1] = v, ("rot_2", Property::Float(v)) => self.rotation[2] = v, @@ -65,5 +71,13 @@ pub fn parse_ply(mut reader: &mut dyn BufRead) -> Result, Error> { } } + for gaussian in &mut cloud { + let mean_scale = (gaussian.scale.x + gaussian.scale.y + gaussian.scale.z) / 3.0; + gaussian.scale = gaussian.scale + .max(Vec3::splat(mean_scale - MAX_SIZE_VARIANCE)) + .min(Vec3::splat(mean_scale + MAX_SIZE_VARIANCE)) + .exp(); + } + Ok(cloud) } diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index a27bedbf..709e3041 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -17,11 +17,12 @@ struct GaussianOutput { @location(0) @interpolate(flat) color: vec4, @location(1) @interpolate(flat) conic: vec3, @location(2) @interpolate(linear) uv: vec2, + @location(3) @interpolate(linear) major_minor: vec2, }; struct GaussianUniforms { + global_transform: mat4x4, global_scale: f32, - transform: f32, }; @@ -36,12 +37,7 @@ struct GaussianUniforms { // https://github.com/cvlab-epfl/gaussian-splatting-web/blob/905b3c0fb8961e42c79ef97e64609e82383ca1c2/src/shaders.ts#L185 // TODO: precompute fn compute_cov3d(scale: vec3, rot: vec4) -> array { - let modifier = uniforms.global_scale; - let S = mat3x3( - scale.x * modifier, 0.0, 0.0, - 0.0, scale.y * modifier, 0.0, - 0.0, 0.0, scale.z * modifier, - ); + let S = scale * uniforms.global_scale; let r = rot.x; let x = rot.y; @@ -49,12 +45,25 @@ fn compute_cov3d(scale: vec3, rot: vec4) -> array { let z = rot.w; let R = mat3x3( - 1.0 - 2.0 * (y * y + z * z), 2.0 * (x * y - r * z), 2.0 * (x * z + r * y), - 2.0 * (x * y + r * z), 1.0 - 2.0 * (x * x + z * z), 2.0 * (y * z - r * x), - 2.0 * (x * z - r * y), 2.0 * (y * z + r * x), 1.0 - 2.0 * (x * x + y * y), + 1.0 - 2.0 * (y * y + z * z), + 2.0 * (x * y - r * z), + 2.0 * (x * z + r * y), + + 2.0 * (x * y + r * z), + 1.0 - 2.0 * (x * x + z * z), + 2.0 * (y * z - r * x), + + 2.0 * (x * z - r * y), + 2.0 * (y * z + r * x), + 1.0 - 2.0 * (x * x + y * y), + ); + + let M = mat3x3( + S[0] * R.x, + S[1] * R.y, + S[2] * R.z, ); - let M = S * R; let Sigma = transpose(M) * M; return array( @@ -69,7 +78,13 @@ fn compute_cov3d(scale: vec3, rot: vec4) -> array { fn compute_cov2d(position: vec3, scale: vec3, rot: vec4) -> vec3 { let cov3d = compute_cov3d(scale, rot); + let Vrk = mat3x3( + cov3d[0], cov3d[1], cov3d[2], + cov3d[1], cov3d[3], cov3d[4], + cov3d[2], cov3d[4], cov3d[5], + ); + // TODO: resolve metal vs directx differences var t = view.inverse_view * vec4(position, 1.0); let focal_x = 500.0; @@ -83,172 +98,110 @@ fn compute_cov2d(position: vec3, scale: vec3, rot: vec4) -> vec3< t.x = min(limx, max(-limx, txtz)) * t.z; t.y = min(limy, max(-limy, tytz)) * t.z; - let J = mat4x4( - focal_x / t.z, 0.0, -(focal_x * t.x) / (t.z * t.z), 0.0, - 0.0, focal_y / t.z, -(focal_y * t.y) / (t.z * t.z), 0.0, - 0.0, 0.0, 0.0, 0.0, - 0.0, 0.0, 0.0, 0.0, - ); + let J = mat3x3( + focal_x / t.z, + 0.0, + -(focal_x * t.x) / (t.z * t.z), - let W = transpose(view.inverse_view); + 0.0, + -focal_y / t.z, + (focal_y * t.y) / (t.z * t.z), - let T = W * J; + 0.0, 0.0, 0.0, + ); - let Vrk = mat4x4( - cov3d[0], cov3d[1], cov3d[2], 0.0, - cov3d[1], cov3d[3], cov3d[4], 0.0, - cov3d[2], cov3d[4], cov3d[5], 0.0, - 0.0, 0.0, 0.0, 0.0, + let W = transpose( + mat3x3( + view.inverse_view.x.xyz, + view.inverse_view.y.xyz, + view.inverse_view.z.xyz, + ) ); - var cov = transpose(T) * transpose(Vrk) * T; + let T = W * J; - // Apply low-pass filter: every Gaussian should be at least - // one pixel wide/high. Discard 3rd row and column. - // cov[0][0] += 0.3; - // cov[1][1] += 0.3; + var cov = transpose(T) * transpose(Vrk) * T; + cov[0][0] += 0.3f; + cov[1][1] += 0.3f; return vec3(cov[0][0], cov[0][1], cov[1][1]); } - -// https://github.com/Lichtso/splatter/blob/c6b7a3894c25578cd29c9761619e4f194449e389/src/shaders.wgsl#L125-L169 -fn quat_to_mat(p: vec4) -> mat3x3 { - var q = p * sqrt(2.0); - var yy = q.y * q.y; - var yz = q.y * q.z; - var yw = q.y * q.w; - var yx = q.y * q.x; - var zz = q.z * q.z; - var zw = q.z * q.w; - var zx = q.z * q.x; - var ww = q.w * q.w; - var wx = q.w * q.x; - return mat3x3( - 1.0 - zz - ww, yz + wx, yw - zx, - yz - wx, 1.0 - yy - ww, zw + yx, - yw + zx, zw - yx, 1.0 - yy - zz, - ); +fn world_to_clip(world_pos: vec3) -> vec4 { + let homogenous_pos = view.view_proj * vec4(world_pos, 1.0); + return homogenous_pos / homogenous_pos.w; } -fn projected_covariance_of_ellipsoid(scale: vec3, rotation: vec4, translation: vec3) -> mat3x3 { - let camera_matrix = mat3x3( - view.view.x.xyz, - view.view.y.xyz, - view.view.z.xyz - ); - var transform = quat_to_mat(rotation); - transform.x *= scale.x; - transform.y *= scale.y; - transform.z *= scale.z; - - // 3D Covariance - var view_pos = view.view * vec4(translation, 1.0); - view_pos.x = clamp(view_pos.x / view_pos.z, -1.0, 1.0) * view_pos.z; - view_pos.y = clamp(view_pos.y / view_pos.z, -1.0, 1.0) * view_pos.z; - let T = transpose(transform) * camera_matrix * mat3x3( - 1.0 / view_pos.z, 0.0, -view_pos.x / (view_pos.z * view_pos.z), - 0.0, 1.0 / view_pos.z, -view_pos.y / (view_pos.z * view_pos.z), - 0.0, 0.0, 0.0, - ); - let covariance_matrix = transpose(T) * T; - - return covariance_matrix; +fn in_frustum(clip_space_pos: vec3) -> bool { + return abs(clip_space_pos.x) < 1.1 + && abs(clip_space_pos.y) < 1.1 + && abs(clip_space_pos.z - 0.5) < 0.5; } -fn projected_contour_of_ellipsoid(scale: vec3, rotation: vec4, translation: vec3) -> mat3x3 { - let camera_matrix = mat3x3( - view.inverse_view.x.xyz, - view.inverse_view.y.xyz, - view.inverse_view.z.xyz - ); - - var transform = quat_to_mat(rotation); - transform.x /= scale.x; - transform.y /= scale.y; - transform.z /= scale.z; - let ray_origin = view.world_position - translation; - let local_ray_origin = ray_origin * transform; - let local_ray_origin_squared = local_ray_origin * local_ray_origin; +fn get_bounding_box_corner( + cov2d: vec3, + direction: vec2, +) -> vec4 { + // return vec4(offset, uv); - let diagonal = 1.0 - local_ray_origin_squared.yxx - local_ray_origin_squared.zzy; - let triangle = local_ray_origin.yxx * local_ray_origin.zzy; + let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y; - let A = mat3x3( - diagonal.x, triangle.z, triangle.y, - triangle.z, diagonal.y, triangle.x, - triangle.y, triangle.x, diagonal.z, + 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); + +#ifdef USE_AABB + // 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, ); - transform = transpose(camera_matrix) * transform; - let M = transform * A * transpose(transform); - - return M; -} - -fn extract_translation_of_ellipse(M: mat3x3) -> vec2 { - let discriminant = M.x.x * M.y.y - M.x.y * M.x.y; - let inverse_discriminant = 1.0 / discriminant; - return vec2( - M.x.y * M.y.z - M.y.y * M.x.z, - M.x.y * M.x.z - M.x.x * M.y.z, - ) * inverse_discriminant; -} - -fn extract_rotation_of_ellipse(M: mat3x3) -> vec2 { - let a = (M.x.x - M.y.y) * (M.x.x - M.y.y); - let b = a + 4.0 * M.x.y * M.x.y; - let c = 0.5 * sqrt(a / b); - var j = sqrt(0.5 - c); - var k = -sqrt(0.5 + c) * sign(M.x.y) * sign(M.x.x - M.y.y); - if(M.x.y < 0.0 || M.x.x - M.y.y < 0.0) { - k = -k; - j = -j; - } - if(M.x.x - M.y.y < 0.0) { - let t = j; - j = -k; - k = t; + return vec4( + 2.0 * radius_ndc * direction, + radius_px * direction, + ); +#endif + +#ifdef USE_OBB + // 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) { + return vec4( + vec2( + direction.x * (x_axis_length + y_axis_length) * 0.5, + direction.y * x_axis_length + ) / view.viewport.zw, + direction * x_axis_length + ); } - return vec2(j, k); -} - -fn extract_scale_of_ellipse(M: mat3x3, translation: vec2, rotation: vec2) -> vec2 { - let d = 2.0 * M.x.y * rotation.x * rotation.y; - let e = M.z.z - (M.x.x * translation.x * translation.x + M.y.y * translation.y * translation.y + 2.0 * M.x.y * translation.x * translation.y); - let semi_major_axis = sqrt(abs(e / (M.x.x * rotation.y * rotation.y + M.y.y * rotation.x * rotation.x - d))); - let semi_minor_axis = sqrt(abs(e / (M.x.x * rotation.x * rotation.x + M.y.y * rotation.y * rotation.y + d))); - - return vec2(semi_major_axis, semi_minor_axis); -} - -fn extract_scale_of_covariance(M: mat3x3) -> vec2 { - let a = (M.x.x - M.y.y) * (M.x.x - M.y.y); - let b = sqrt(a + 4.0 * M.x.y * M.x.y); - let semi_major_axis = sqrt((M.x.x + M.y.y + b) * 0.5); - let semi_minor_axis = sqrt((M.x.x + M.y.y - b) * 0.5); - return vec2(semi_major_axis, semi_minor_axis); -} -fn world_to_clip(world_pos: vec3) -> vec4 { - let homogenous_pos = view.view_proj * vec4(world_pos, 1.0); - return vec4(homogenous_pos.xyz, 1.0) / (homogenous_pos.w + 0.0000001); -} - -fn in_frustum(clip_space_pos: vec3) -> bool { - return abs(clip_space_pos.x) < 1.1 - && abs(clip_space_pos.y) < 1.1 - && abs(clip_space_pos.z - 0.5) < 0.5; -} + let eigvec1 = normalize(vec2( + cov2d.y, + lambda1 - cov2d.x + )); + let eigvec2 = vec2(-eigvec1.y, eigvec1.x); + let rotation_matrix = mat2x2( + eigvec1.x, eigvec2.x, + eigvec1.y, eigvec2.y + ); -fn view_dimensions(projection: mat4x4) -> vec2 { - let near = projection[2][3] / (projection[2][2] + 1.0); - let right = near / projection[0][0]; - let top = near / projection[1][1]; + let scaled_vertex = vec2( + direction.x * x_axis_length, + direction.y * y_axis_length + ); - return vec2(2.0 * right, 2.0 * top); + return vec4( + rotation_matrix * (scaled_vertex / view.viewport.zw), + scaled_vertex + ); +#endif } @@ -259,8 +212,10 @@ fn vs_points( ) -> GaussianOutput { var output: GaussianOutput; let point = points[instance_index]; + let transformed_position = (uniforms.global_transform * vec4(point.position, 1.0)).xyz; - if (!in_frustum(world_to_clip(point.position).xyz)) { + let projected_position = world_to_clip(transformed_position); + if (!in_frustum(projected_position.xyz)) { output.color = vec4(0.0, 0.0, 0.0, 0.0); return output; } @@ -275,17 +230,17 @@ fn vs_points( let quad_index = vertex_index % 4u; let quad_offset = quad_vertices[quad_index]; - let ray_direction = normalize(point.position - view.world_position); + let ray_direction = normalize(transformed_position - view.world_position); output.color = vec4( spherical_harmonics_lookup(ray_direction, point.sh), point.opacity ); - let cov2d = compute_cov2d(point.position, point.scale, point.rotation); + let cov2d = compute_cov2d(transformed_position, point.scale, 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( cov2d.z * det_inv, -cov2d.y * det_inv, @@ -293,80 +248,25 @@ fn vs_points( ); output.conic = conic; - let mid = 0.5 * (cov2d.x + cov2d.z); - let lambda_1 = mid + sqrt(max(0.1, mid * mid - det)); - let lambda_2 = mid - sqrt(max(0.1, mid * mid - det)); - let radius_px = 3.5 * sqrt(max(lambda_1, lambda_2)); - let radius_ndc = vec2( - radius_px / f32(view.viewport.z), - radius_px / f32(view.viewport.w), + let bb = get_bounding_box_corner( + cov2d, + quad_offset, ); - output.uv = radius_px * quad_offset; - - var projected_position = view.view_proj * vec4(point.position, 1.0); - projected_position = projected_position / projected_position.w; - + output.uv = (quad_offset + vec2(1.0)) * 0.5; + output.major_minor = bb.zw; output.position = vec4( - projected_position.xy + 2.0 * radius_ndc * quad_offset, - projected_position.zw, + projected_position.xy + bb.xy, + projected_position.zw ); - // let M = projected_contour_of_ellipsoid( - // point.scale * uniforms.global_scale, - // point.rotation, - // point.position, - // ); - // let translation = extract_translation_of_ellipse(M); - // let rotation = extract_rotation_of_ellipse(M); - // //let semi_axes = extract_scale_of_ellipse(M, translation, rotation); - - // let covariance = projected_covariance_of_ellipsoid( - // point.scale * uniforms.global_scale, - // point.rotation, - // point.position - // ); - // let semi_axes = extract_scale_of_covariance(covariance); - - // let view_dimensions = view_dimensions(view.projection); - // let ellipse_size_bias = 0.2 * view_dimensions.x / f32(view.viewport.z); - - // let transformation = mat3x2( - // vec2(rotation.y, -rotation.x) * (ellipse_size_bias + semi_axes.x), - // vec2(rotation.x, rotation.y) * (ellipse_size_bias + semi_axes.y), - // translation, - // ); - - // let T = mat3x3( - // vec3(transformation.x, 0.0), - // vec3(transformation.y, 0.0), - // vec3(transformation.z, 1.0), - // ); - - // let ellipse_margin = 3.3; // should be 2.0 - // output.uv = quad_offset * ellipse_margin; - // output.position = vec4( - // (T * vec3(output.uv, 1.0)).xy / view_dimensions, - // 0.0, - // 1.0, - // ); - return output; } @fragment fn fs_main(input: GaussianOutput) -> @location(0) vec4 { - // let power = dot(input.uv, input.uv); - // let alpha = input.color.a * exp(-0.5 * power); - - // if (alpha < 1.0 / 255.0) { - // discard; - // } - - // return vec4(input.color.rgb * alpha, alpha); - - - let d = -input.uv; + // TODO: draw gaussian without conic (OBB) + let d = -input.major_minor; let conic = input.conic; let power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y; @@ -374,6 +274,17 @@ fn fs_main(input: GaussianOutput) -> @location(0) vec4 { discard; } +#ifdef VISUALIZE_BOUNDING_BOX + let uv = input.uv; + let edge_width = 0.08; + if ( + (uv.x < edge_width || uv.x > 1.0 - edge_width) || + (uv.y < edge_width || uv.y > 1.0 - edge_width) + ) { + return vec4(0.3, 1.0, 0.1, 1.0); + } +#endif + let alpha = min(0.99, input.color.a * exp(power)); return vec4( input.color.rgb * alpha, diff --git a/src/render/mod.rs b/src/render/mod.rs index ac402b27..4cdf8e32 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -120,6 +120,7 @@ impl Plugin for RenderPipelinePlugin { #[derive(Bundle)] pub struct GpuGaussianSplattingBundle { + pub settings: GaussianCloudSettings, pub settings_uniform: GaussianCloudUniform, pub verticies: Handle, } @@ -168,16 +169,18 @@ fn queue_gaussians( gaussian_splatting_bundles: Query<( Entity, &Handle, + &GaussianCloudSettings, )>, mut views: Query<(&ExtractedView, &mut RenderPhase)>, ) { let draw_custom = transparent_3d_draw_functions.read().id::(); for (_view, mut transparent_phase) in &mut views { - for (entity, verticies) in &gaussian_splatting_bundles { - if let Some(_cloud) = gaussian_clouds.get(verticies) { + for (entity, cloud, settings) in &gaussian_splatting_bundles { + if let Some(_cloud) = gaussian_clouds.get(cloud) { let key = GaussianCloudPipelineKey { - + aabb: settings.aabb, + visualize_bounding_box: settings.visualize_bounding_box, }; let pipeline = pipelines.specialize(&pipeline_cache, &custom_pipeline, key); @@ -277,17 +280,30 @@ impl FromWorld for GaussianCloudPipeline { #[derive(PartialEq, Eq, Hash, Clone, Copy)] pub struct GaussianCloudPipelineKey { - + pub aabb: bool, + pub visualize_bounding_box: bool, } impl SpecializedRenderPipeline for GaussianCloudPipeline { type Key = GaussianCloudPipelineKey; - fn specialize(&self, _key: Self::Key) -> RenderPipelineDescriptor { - let shader_defs = vec![ + fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor { + let mut shader_defs = vec![ ShaderDefVal::UInt("MAX_SH_COEFF_COUNT".into(), MAX_SH_COEFF_COUNT as u32), ]; + if key.aabb { + shader_defs.push("USE_AABB".into()); + } + + if !key.aabb { + shader_defs.push("USE_OBB".into()); + } + + if key.visualize_bounding_box { + shader_defs.push("VISUALIZE_BOUNDING_BOX".into()); + } + RenderPipelineDescriptor { label: Some("gaussian cloud pipeline".into()), layout: vec![ @@ -367,8 +383,8 @@ type DrawGaussians = ( #[derive(Component, ShaderType, Clone)] pub struct GaussianCloudUniform { - pub global_scale: f32, pub transform: Mat4, + pub global_scale: f32, } pub fn extract_gaussians( @@ -378,24 +394,25 @@ pub fn extract_gaussians( Query<( Entity, // &ComputedVisibility, - &GaussianCloudSettings, &Handle, + &GaussianCloudSettings, )>, >, ) { let mut commands_list = Vec::with_capacity(*prev_commands_len); // let visible_gaussians = gaussians_query.iter().filter(|(_, vis, ..)| vis.is_visible()); - for (entity, settings, verticies) in gaussians_query.iter() { + for (entity, verticies, settings) in gaussians_query.iter() { let settings_uniform = GaussianCloudUniform { - global_scale: settings.global_scale, transform: settings.global_transform.compute_matrix(), + global_scale: settings.global_scale, }; commands_list.push(( entity, GpuGaussianSplattingBundle { + settings: settings.clone(), settings_uniform, - verticies: verticies.clone_weak(), + verticies: verticies.clone(), }, )); } @@ -433,27 +450,31 @@ pub fn queue_gaussian_bind_group( assert!(model.size() == std::mem::size_of::() as u64); + groups.base_bind_group = Some(render_device.create_bind_group(&BindGroupDescriptor { + entries: &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::Buffer(BufferBinding { + buffer: model, + offset: 0, + size: BufferSize::new(model.size()), + }), + }, + ], + layout: &gaussian_cloud_pipeline.gaussian_uniform_layout, + label: Some("gaussian_uniform_bind_group"), + })); + for (entity, cloud_handle) in gaussian_clouds.iter() { if asset_server.get_load_state(cloud_handle) == LoadState::Loading { continue; } - let cloud = gaussian_cloud_res.get(cloud_handle).unwrap(); + if !gaussian_cloud_res.contains_key(cloud_handle) { + continue; + } - groups.base_bind_group = Some(render_device.create_bind_group(&BindGroupDescriptor { - entries: &[ - BindGroupEntry { - binding: 0, - resource: BindingResource::Buffer(BufferBinding { - buffer: model, - offset: 0, - size: BufferSize::new(model.size()), - }), - }, - ], - layout: &gaussian_cloud_pipeline.gaussian_uniform_layout, - label: Some("gaussian_uniform_bind_group"), - })); + let cloud = gaussian_cloud_res.get(cloud_handle).unwrap(); commands.entity(entity).insert(GaussianCloudBindGroup { bind_group: render_device.create_bind_group(&BindGroupDescriptor { @@ -624,9 +645,10 @@ impl RenderCommand

for DrawGaussianInstanced { GpuBufferInfo::NonIndexed => { pass.draw(0..4, 0..gpu_gaussian_cloud.count as u32); } - // TODO: add support for indirect draw and match over sort methods } RenderCommandResult::Success } + + } diff --git a/tools/ply_to_gcloud.rs b/tools/ply_to_gcloud.rs new file mode 100644 index 00000000..ae167c8e --- /dev/null +++ b/tools/ply_to_gcloud.rs @@ -0,0 +1,37 @@ +use bincode2::serialize_into; +use flate2::{ + Compression, + write::GzEncoder, +}; + +use bevy_gaussian_splatting::{ + GaussianCloud, + ply::parse_ply, +}; + + +fn main() { + let filename = std::env::args().nth(1).expect("no filename given"); + + println!("converting {}", filename); + + // filepath to BufRead + let file = std::fs::File::open(&filename).expect("failed to open file"); + let mut reader = std::io::BufReader::new(file); + + let cloud = GaussianCloud(parse_ply(&mut reader).expect("failed to parse ply file")); + + // write cloud to .gcloud file (remove .ply) + let base_filename = filename.split('.').next().expect("no extension").to_string(); + let gcloud_filename = base_filename + ".gcloud"; + // let gcloud_file = std::fs::File::create(&gcloud_filename).expect("failed to create file"); + // let mut writer = std::io::BufWriter::new(gcloud_file); + + // serialize_into(&mut writer, &cloud).expect("failed to encode cloud"); + + // write gloud.gz + let gz_file = std::fs::File::create(&gcloud_filename).expect("failed to create file"); + let mut gz_writer = std::io::BufWriter::new(gz_file); + let mut gz_encoder = GzEncoder::new(&mut gz_writer, Compression::default()); // TODO: consider switching to fast (or support multiple options), default is a bit slow + serialize_into(&mut gz_encoder, &cloud).expect("failed to encode cloud"); +}