From 60331951298d7bb40274b3ea1a2472ee64f9af0a Mon Sep 17 00:00:00 2001 From: Mitchell Mosure Date: Fri, 15 Dec 2023 17:53:31 -0600 Subject: [PATCH] feat: web (#52) * feat: running in the web * docs: add TODO --- .cargo/config.toml | 6 +- .github/workflows/test.yml | 3 + .gitignore | 3 + .../bevy_gaussian_splatting.code-workspace | 4 +- Cargo.toml | 21 +- README.md | 10 +- src/morph/mod.rs | 372 +----------------- src/morph/particle.rs | 371 ++++++++++++++++- src/morph/particle.wgsl | 8 +- src/render/bindings.wgsl | 5 + src/render/mod.rs | 11 +- src/sort/mod.rs | 13 + src/sort/std.rs | 102 +++++ src/utils.rs | 30 +- tests/gpu/gaussian.rs | 2 + tests/gpu/radix.rs | 10 +- tools/compare_aabb_obb.rs | 2 + viewer/viewer.rs | 100 +++-- www/README.md | 15 + www/index.html | 23 ++ 20 files changed, 685 insertions(+), 426 deletions(-) create mode 100644 src/sort/std.rs create mode 100644 www/README.md create mode 100644 www/index.html diff --git a/.cargo/config.toml b/.cargo/config.toml index 805f8358..ad439478 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -2,7 +2,11 @@ [target.wasm32-unknown-unknown] runner = "wasm-server-runner" -rustflags = ["--cfg=web_sys_unstable_apis"] +rustflags = [ + "--cfg=web_sys_unstable_apis", + # "-C", + # "target-feature=+atomics,+bulk-memory,+mutable-globals", # for wasm-bindgen-rayon +] # fix spurious network error on windows diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 693d14a8..99994c24 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -50,3 +50,6 @@ jobs: # - name: radix sort test # run: cargo run --bin test_radix --features="debug_gpu" + + + # TODO: test wasm build, deploy, and run diff --git a/.gitignore b/.gitignore index 45d282a6..19d17af9 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ debug/ target/ +out/ Cargo.lock **/*.rs.bk *.pdb @@ -10,3 +11,5 @@ Cargo.lock *.gcloud .DS_Store + +www/assets/ diff --git a/.vscode/bevy_gaussian_splatting.code-workspace b/.vscode/bevy_gaussian_splatting.code-workspace index b355fe5c..fbf610aa 100644 --- a/.vscode/bevy_gaussian_splatting.code-workspace +++ b/.vscode/bevy_gaussian_splatting.code-workspace @@ -7,5 +7,7 @@ "path": "../../bevy" } ], - "settings": {} + "settings": { + "liveServer.settings.multiRootWorkspaceName": "bevy_gaussian_splatting" + } } \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index 99e1d334..eaf9d398 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,8 +36,10 @@ default-run = "viewer" default = [ "io_flexbuffers", "io_ply", + "morph_particles", "sort_radix", "sort_rayon", + "sort_std", "tooling", "viewer", ] @@ -48,11 +50,13 @@ io_bincode2 = ["bincode2", "flate2"] io_flexbuffers = ["flexbuffers"] io_ply = ["ply-rs"] +morph_particles = [] + sort_radix = [] sort_rayon = [ "rayon", - "wasm-bindgen-rayon", ] +sort_std = [] tooling = [ "byte-unit", @@ -64,7 +68,6 @@ viewer = [ ] - [dependencies] bevy-inspector-egui = { version = "0.21", optional = true } bevy_panorbit_camera = { version = "0.9", optional = true } @@ -84,29 +87,27 @@ wgpu = "0.17.1" [target.'cfg(target_arch = "wasm32")'.dependencies] console_error_panic_hook = "0.1" wasm-bindgen = "0.2" -wasm-bindgen-rayon = { version = "1.0", optional = true } # TODO: use minimal bevy features [dependencies.bevy] version = "0.12" -default-features = true - -[dev-dependencies.bevy] -version = "0.12" -default-features = true +default-features = false features = [ - 'debug_glam_assert', + "bevy_asset", + "bevy_core_pipeline", + "bevy_render", + "bevy_winit", ] - [dependencies.web-sys] version = "0.3.4" features = [ 'Document', 'Element', 'HtmlElement', + 'Location', 'Node', 'Window', ] diff --git a/README.md b/README.md index 191c0316..e203f62c 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,9 @@ download [cactus.gcloud](https://mitchell.mosure.me/cactus.gcloud) - [X] gcloud and ply asset loaders - [X] bevy gaussian cloud render pipeline - [X] gaussian cloud particle effects +- [X] wasm support /w [live demo](https://mosure.github.io/bevy_gaussian_splatting/index.html?arg1=icecream.gcloud) - [ ] 4D gaussian cloud wavelet compression - [ ] accelerated spatial queries -- [ ] wasm support /w [live demo](https://mosure.github.io/bevy_gaussian_splatting) - [ ] temporal depth sorting - [ ] f16 and f32 gcloud support - [ ] skeletons @@ -69,18 +69,12 @@ fn setup_gaussian_cloud( - [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 - -to build wasm run: -- `cargo build --target wasm32-unknown-unknown --release` -- `wasm-bindgen --out-dir ./out/ --target web ./target/` - ## compatible bevy versions | `bevy_gaussian_splatting` | `bevy` | | :-- | :-- | -| `0.4` | `0.12` | +| `0.4 - 0.5` | `0.12` | | `0.1 - 0.3` | `0.11` | diff --git a/src/morph/mod.rs b/src/morph/mod.rs index b8e0d2d3..fd2c87f8 100644 --- a/src/morph/mod.rs +++ b/src/morph/mod.rs @@ -1,379 +1,17 @@ -use bevy::{ - prelude::*, - asset::{ - load_internal_asset, - LoadState, - }, - core_pipeline::core_3d::CORE_3D, - ecs::system::{ - lifetimeless::SRes, - SystemParamItem, - }, - render::{ - Extract, - render_asset::{ - PrepareAssetError, - RenderAsset, - RenderAssets, - RenderAssetPlugin, - }, - render_resource::*, - renderer::{ - RenderContext, - RenderDevice, - }, - render_graph::{ - Node, - NodeRunError, - RenderGraphApp, - RenderGraphContext, - }, - Render, - RenderApp, - RenderSet, - view::ViewUniformOffset, - }, -}; +use bevy::prelude::*; -use crate::render::{ - GaussianCloudBindGroup, - GaussianCloudPipeline, - GaussianUniformBindGroups, - GaussianViewBindGroup, - shader_defs, -}; - -pub use particle::{ - ParticleBehavior, - ParticleBehaviors, - random_particle_behaviors, -}; +#[cfg(feature = "morph_particles")] pub mod particle; -const PARTICLE_SHADER_HANDLE: Handle = Handle::weak_from_u128(234553453455); -pub mod node { - pub const MORPH: &str = "gaussian_cloud_morph"; -} - - #[derive(Default)] pub struct MorphPlugin; impl Plugin for MorphPlugin { + #[allow(unused)] fn build(&self, app: &mut App) { - load_internal_asset!( - app, - PARTICLE_SHADER_HANDLE, - "particle.wgsl", - Shader::from_wgsl - ); - - app.register_type::(); - app.init_asset::(); - app.register_asset_reflect::(); - app.add_plugins(RenderAssetPlugin::::default()); - - if let Ok(render_app) = app.get_sub_app_mut(RenderApp) { - render_app - .add_render_graph_node::( - CORE_3D, - node::MORPH, - ) - .add_render_graph_edge( - CORE_3D, - node::MORPH, - bevy::core_pipeline::core_3d::graph::node::PREPASS, - ); - - render_app - .add_systems(ExtractSchedule, extract_particle_behaviors) - .add_systems( - Render, - ( - queue_morph_bind_group.in_set(RenderSet::QueueMeshes), - ), - ); - } - } - - fn finish(&self, app: &mut App) { - if let Ok(render_app) = app.get_sub_app_mut(RenderApp) { - render_app - .init_resource::(); - } - } -} - - -pub fn extract_particle_behaviors( - mut commands: Commands, - mut prev_commands_len: Local, - gaussians_query: Extract< - Query<( - Entity, - &Handle, - )>, - >, -) { - let mut commands_list = Vec::with_capacity(*prev_commands_len); - - for (entity, behaviors) in gaussians_query.iter() { - commands_list.push(( - entity, - behaviors.clone(), - )); - } - *prev_commands_len = commands_list.len(); - commands.insert_or_spawn_batch(commands_list); -} - - -#[derive(Debug, Clone)] -pub struct GpuMorphBuffers { - pub morph_count: u32, - pub particle_behavior_buffer: Buffer, -} - -impl RenderAsset for ParticleBehaviors { - type ExtractedAsset = ParticleBehaviors; - type PreparedAsset = GpuMorphBuffers; - type Param = SRes; - - fn extract_asset(&self) -> Self::ExtractedAsset { - self.clone() - } - - fn prepare_asset( - particle_behaviors: Self::ExtractedAsset, - render_device: &mut SystemParamItem, - ) -> Result> { - let morph_count = particle_behaviors.0.len() as u32; - - let particle_behavior_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor { - label: Some("particle behavior buffer"), - contents: bytemuck::cast_slice( - particle_behaviors.0.as_slice() - ), - usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, - }); - - Ok(GpuMorphBuffers { - morph_count, - particle_behavior_buffer, - }) - } -} - - -#[derive(Resource)] -pub struct MorphPipeline { - pub morph_layout: BindGroupLayout, - pub particle_behavior_pipeline: CachedComputePipelineId, -} - -impl FromWorld for MorphPipeline { - fn from_world(render_world: &mut World) -> Self { - let render_device = render_world.resource::(); - let gaussian_cloud_pipeline = render_world.resource::(); - - let morph_layout = render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { - label: Some("gaussian_cloud_morph_layout"), - entries: &[ - BindGroupLayoutEntry { - binding: 7, - visibility: ShaderStages::COMPUTE, - ty: BindingType::Buffer { - ty: BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: BufferSize::new(std::mem::size_of::() as u64), - }, - count: None, - }, - ], - }); - - let shader_defs = shader_defs(false, false); - let pipeline_cache = render_world.resource::(); - - let particle_behavior_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { - label: Some("particle_behavior_pipeline".into()), - layout: vec![ - gaussian_cloud_pipeline.view_layout.clone(), - gaussian_cloud_pipeline.gaussian_uniform_layout.clone(), - gaussian_cloud_pipeline.gaussian_cloud_layout.clone(), - morph_layout.clone(), - ], - push_constant_ranges: vec![], - shader: PARTICLE_SHADER_HANDLE, - shader_defs: shader_defs.clone(), - entry_point: "apply_particle_behaviors".into(), - }); - - MorphPipeline { - morph_layout, - particle_behavior_pipeline, - } - } -} - - - -#[derive(Component)] -pub struct MorphBindGroup { - pub morph_bindgroup: BindGroup, -} - -pub fn queue_morph_bind_group( - mut commands: Commands, - morph_pipeline: Res, - render_device: Res, - asset_server: Res, - particle_behaviors_res: Res>, - particle_behaviors: Query<( - Entity, - &Handle, - )>, -) { - for (entity, behaviors_handle) in particle_behaviors.iter() { - if Some(LoadState::Loading) == asset_server.get_load_state(behaviors_handle) { - continue; - } - - if particle_behaviors_res.get(behaviors_handle).is_none() { - continue; - } - - let behaviors = particle_behaviors_res.get(behaviors_handle).unwrap(); - - let morph_bindgroup = render_device.create_bind_group( - "morph_bind_group", - &morph_pipeline.morph_layout, - &[ - BindGroupEntry { - binding: 7, - resource: BindingResource::Buffer(BufferBinding { - buffer: &behaviors.particle_behavior_buffer, - offset: 0, - size: BufferSize::new(behaviors.particle_behavior_buffer.size()), - }), - }, - ], - ); - - commands.entity(entity).insert(MorphBindGroup { - morph_bindgroup, - }); - } -} - - - -pub struct MorphNode { - gaussian_clouds: QueryState<( - &'static GaussianCloudBindGroup, - &'static Handle, - &'static MorphBindGroup, - )>, - initialized: bool, - view_bind_group: QueryState<( - &'static GaussianViewBindGroup, - &'static ViewUniformOffset, - )>, -} - - -impl FromWorld for MorphNode { - fn from_world(world: &mut World) -> Self { - Self { - gaussian_clouds: world.query(), - initialized: false, - view_bind_group: world.query(), - } - } -} - -impl Node for MorphNode { - fn update(&mut self, world: &mut World) { - let pipeline = world.resource::(); - let pipeline_cache = world.resource::(); - - if !self.initialized { - if let CachedPipelineState::Ok(_) = - pipeline_cache.get_compute_pipeline_state(pipeline.particle_behavior_pipeline) - { - self.initialized = true; - } - - if !self.initialized { - return; - } - } - - - self.gaussian_clouds.update_archetypes(world); - self.view_bind_group.update_archetypes(world); - } - - fn run( - &self, - _graph: &mut RenderGraphContext, - render_context: &mut RenderContext, - world: &World, - ) -> Result<(), NodeRunError> { - if !self.initialized { - return Ok(()); - } - - let pipeline_cache = world.resource::(); - let pipeline = world.resource::(); - - let command_encoder = render_context.command_encoder(); - - for ( - view_bind_group, - view_uniform_offset, - ) in self.view_bind_group.iter_manual(world) { - for ( - cloud_bind_group, - behaviors_handle, - morph_bind_group, - ) in self.gaussian_clouds.iter_manual(world) { - let behaviors = world.get_resource::>().unwrap().get(behaviors_handle).unwrap(); - let gaussian_uniforms = world.resource::(); - - { - let mut pass = command_encoder.begin_compute_pass(&ComputePassDescriptor::default()); - - pass.set_bind_group( - 0, - &view_bind_group.value, - &[view_uniform_offset.offset], - ); - pass.set_bind_group( - 1, - gaussian_uniforms.base_bind_group.as_ref().unwrap(), - &[0], - ); - pass.set_bind_group( - 2, - &cloud_bind_group.cloud_bind_group, - &[] - ); - pass.set_bind_group( - 3, - &morph_bind_group.morph_bindgroup, - &[], - ); - - let particle_behavior = pipeline_cache.get_compute_pipeline(pipeline.particle_behavior_pipeline).unwrap(); - pass.set_pipeline(particle_behavior); - pass.dispatch_workgroups(behaviors.morph_count / 32, 32, 1); - } - } - } - - Ok(()) + #[cfg(feature = "morph_particles")] + app.add_plugins(particle::ParticleBehaviorPlugin); } } diff --git a/src/morph/particle.rs b/src/morph/particle.rs index 1143dede..e728e042 100644 --- a/src/morph/particle.rs +++ b/src/morph/particle.rs @@ -6,8 +6,40 @@ use std::marker::Copy; use bevy::{ prelude::*, + asset::{ + load_internal_asset, + LoadState, + }, + core_pipeline::core_3d::CORE_3D, + ecs::system::{ + lifetimeless::SRes, + SystemParamItem, + }, reflect::TypeUuid, - render::render_resource::ShaderType, + render::{ + Extract, + render_asset::{ + PrepareAssetError, + RenderAsset, + RenderAssets, + RenderAssetPlugin, + }, + render_resource::*, + renderer::{ + RenderContext, + RenderDevice, + }, + render_graph::{ + Node, + NodeRunError, + RenderGraphApp, + RenderGraphContext, + }, + Render, + RenderApp, + RenderSet, + view::ViewUniformOffset, + }, }; use bytemuck::{ Pod, @@ -18,6 +50,343 @@ use serde::{ Serialize, }; +use crate::render::{ + GaussianCloudBindGroup, + GaussianCloudPipeline, + GaussianUniformBindGroups, + GaussianViewBindGroup, + shader_defs, +}; + + +const PARTICLE_SHADER_HANDLE: Handle = Handle::weak_from_u128(234553453455); +pub mod node { + pub const MORPH: &str = "gaussian_cloud_particle_behavior"; +} + + +#[derive(Default)] +pub struct ParticleBehaviorPlugin; + +impl Plugin for ParticleBehaviorPlugin { + fn build(&self, app: &mut App) { + load_internal_asset!( + app, + PARTICLE_SHADER_HANDLE, + "particle.wgsl", + Shader::from_wgsl + ); + + app.register_type::(); + app.init_asset::(); + app.register_asset_reflect::(); + app.add_plugins(RenderAssetPlugin::::default()); + + if let Ok(render_app) = app.get_sub_app_mut(RenderApp) { + render_app + .add_render_graph_node::( + CORE_3D, + node::MORPH, + ) + .add_render_graph_edge( + CORE_3D, + node::MORPH, + bevy::core_pipeline::core_3d::graph::node::PREPASS, + ); + + render_app + .add_systems(ExtractSchedule, extract_particle_behaviors) + .add_systems( + Render, + ( + queue_particle_behavior_bind_group.in_set(RenderSet::QueueMeshes), + ), + ); + } + } + + fn finish(&self, app: &mut App) { + if let Ok(render_app) = app.get_sub_app_mut(RenderApp) { + render_app + .init_resource::(); + } + } +} + + +pub fn extract_particle_behaviors( + mut commands: Commands, + mut prev_commands_len: Local, + gaussians_query: Extract< + Query<( + Entity, + &Handle, + )>, + >, +) { + let mut commands_list = Vec::with_capacity(*prev_commands_len); + + for (entity, behaviors) in gaussians_query.iter() { + commands_list.push(( + entity, + behaviors.clone(), + )); + } + *prev_commands_len = commands_list.len(); + commands.insert_or_spawn_batch(commands_list); +} + + +#[derive(Debug, Clone)] +pub struct GpuParticleBehaviorBuffers { + pub particle_behavior_count: u32, + pub particle_behavior_buffer: Buffer, +} + +impl RenderAsset for ParticleBehaviors { + type ExtractedAsset = ParticleBehaviors; + type PreparedAsset = GpuParticleBehaviorBuffers; + type Param = SRes; + + fn extract_asset(&self) -> Self::ExtractedAsset { + self.clone() + } + + fn prepare_asset( + particle_behaviors: Self::ExtractedAsset, + render_device: &mut SystemParamItem, + ) -> Result> { + let particle_behavior_count = particle_behaviors.0.len() as u32; + + let particle_behavior_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("particle behavior buffer"), + contents: bytemuck::cast_slice( + particle_behaviors.0.as_slice() + ), + usage: BufferUsages::VERTEX | BufferUsages::COPY_DST | BufferUsages::STORAGE, + }); + + Ok(GpuParticleBehaviorBuffers { + particle_behavior_count, + particle_behavior_buffer, + }) + } +} + + +#[derive(Resource)] +pub struct ParticleBehaviorPipeline { + pub particle_behavior_layout: BindGroupLayout, + pub particle_behavior_pipeline: CachedComputePipelineId, +} + +impl FromWorld for ParticleBehaviorPipeline { + fn from_world(render_world: &mut World) -> Self { + let render_device = render_world.resource::(); + let gaussian_cloud_pipeline = render_world.resource::(); + + let particle_behavior_layout = render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: Some("gaussian_cloud_particle_behavior_layout"), + entries: &[ + BindGroupLayoutEntry { + binding: 7, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + ], + }); + + let shader_defs = shader_defs(false, false); + let pipeline_cache = render_world.resource::(); + + let particle_behavior_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("particle_behavior_pipeline".into()), + layout: vec![ + gaussian_cloud_pipeline.view_layout.clone(), + gaussian_cloud_pipeline.gaussian_uniform_layout.clone(), + gaussian_cloud_pipeline.gaussian_cloud_layout.clone(), + particle_behavior_layout.clone(), + ], + push_constant_ranges: vec![], + shader: PARTICLE_SHADER_HANDLE, + shader_defs: shader_defs.clone(), + entry_point: "apply_particle_behaviors".into(), + }); + + ParticleBehaviorPipeline { + particle_behavior_layout, + particle_behavior_pipeline, + } + } +} + + + +#[derive(Component)] +pub struct ParticleBehaviorBindGroup { + pub particle_behavior_bindgroup: BindGroup, +} + +pub fn queue_particle_behavior_bind_group( + mut commands: Commands, + particle_behavior_pipeline: Res, + render_device: Res, + asset_server: Res, + particle_behaviors_res: Res>, + particle_behaviors: Query<( + Entity, + &Handle, + )>, +) { + for (entity, behaviors_handle) in particle_behaviors.iter() { + if Some(LoadState::Loading) == asset_server.get_load_state(behaviors_handle) { + continue; + } + + if particle_behaviors_res.get(behaviors_handle).is_none() { + continue; + } + + let behaviors = particle_behaviors_res.get(behaviors_handle).unwrap(); + + let particle_behavior_bindgroup = render_device.create_bind_group( + "particle_behavior_bind_group", + &particle_behavior_pipeline.particle_behavior_layout, + &[ + BindGroupEntry { + binding: 7, + resource: BindingResource::Buffer(BufferBinding { + buffer: &behaviors.particle_behavior_buffer, + offset: 0, + size: BufferSize::new(behaviors.particle_behavior_buffer.size()), + }), + }, + ], + ); + + commands.entity(entity).insert(ParticleBehaviorBindGroup { + particle_behavior_bindgroup, + }); + } +} + + + +pub struct ParticleBehaviorNode { + gaussian_clouds: QueryState<( + &'static GaussianCloudBindGroup, + &'static Handle, + &'static ParticleBehaviorBindGroup, + )>, + initialized: bool, + view_bind_group: QueryState<( + &'static GaussianViewBindGroup, + &'static ViewUniformOffset, + )>, +} + + +impl FromWorld for ParticleBehaviorNode { + fn from_world(world: &mut World) -> Self { + Self { + gaussian_clouds: world.query(), + initialized: false, + view_bind_group: world.query(), + } + } +} + +impl Node for ParticleBehaviorNode { + fn update(&mut self, world: &mut World) { + let pipeline = world.resource::(); + let pipeline_cache = world.resource::(); + + if !self.initialized { + if let CachedPipelineState::Ok(_) = + pipeline_cache.get_compute_pipeline_state(pipeline.particle_behavior_pipeline) + { + self.initialized = true; + } + + if !self.initialized { + return; + } + } + + + self.gaussian_clouds.update_archetypes(world); + self.view_bind_group.update_archetypes(world); + } + + fn run( + &self, + _graph: &mut RenderGraphContext, + render_context: &mut RenderContext, + world: &World, + ) -> Result<(), NodeRunError> { + if !self.initialized { + return Ok(()); + } + + let pipeline_cache = world.resource::(); + let pipeline = world.resource::(); + + let command_encoder = render_context.command_encoder(); + + for ( + view_bind_group, + view_uniform_offset, + ) in self.view_bind_group.iter_manual(world) { + for ( + cloud_bind_group, + behaviors_handle, + particle_behavior_bind_group, + ) in self.gaussian_clouds.iter_manual(world) { + let behaviors = world.get_resource::>().unwrap().get(behaviors_handle).unwrap(); + let gaussian_uniforms = world.resource::(); + + { + let mut pass = command_encoder.begin_compute_pass(&ComputePassDescriptor::default()); + + pass.set_bind_group( + 0, + &view_bind_group.value, + &[view_uniform_offset.offset], + ); + pass.set_bind_group( + 1, + gaussian_uniforms.base_bind_group.as_ref().unwrap(), + &[0], + ); + pass.set_bind_group( + 2, + &cloud_bind_group.cloud_bind_group, + &[] + ); + pass.set_bind_group( + 3, + &particle_behavior_bind_group.particle_behavior_bindgroup, + &[], + ); + + let particle_behavior = pipeline_cache.get_compute_pipeline(pipeline.particle_behavior_pipeline).unwrap(); + pass.set_pipeline(particle_behavior); + pass.dispatch_workgroups(behaviors.particle_behavior_count / 32, 32, 1); + } + } + } + + Ok(()) + } +} + + + // TODO: add more particle system functionality (e.g. lifetime, color) #[derive( diff --git a/src/morph/particle.wgsl b/src/morph/particle.wgsl index 82165d5b..002d48bb 100644 --- a/src/morph/particle.wgsl +++ b/src/morph/particle.wgsl @@ -30,10 +30,6 @@ fn apply_particle_behaviors( let behavior_index = gl_GlobalInvocationID.x * 32u + gl_GlobalInvocationID.y; let behavior = particle_behaviors[behavior_index]; - if (behavior.indicies.x < 0) { - return; - } - let point_index = behavior.indicies.x; let point = points[point_index]; @@ -47,6 +43,10 @@ fn apply_particle_behaviors( workgroupBarrier(); + if (behavior.indicies.x < 0) { + return; + } + points[point_index].position = new_position; particle_behaviors[behavior_index].velocity = new_velocity; particle_behaviors[behavior_index].acceleration = new_acceleration; diff --git a/src/render/bindings.wgsl b/src/render/bindings.wgsl index f779a590..8e641b09 100644 --- a/src/render/bindings.wgsl +++ b/src/render/bindings.wgsl @@ -19,7 +19,12 @@ struct Gaussian { @location(2) scale_opacity: vec4, sh: array, }; + +#ifdef READ_WRITE_POINTS @group(2) @binding(0) var points: array; +#else +@group(2) @binding(0) var points: array; +#endif struct DrawIndirect { diff --git a/src/render/mod.rs b/src/render/mod.rs index 5f5974cc..ce2d0238 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -334,6 +334,12 @@ impl FromWorld for GaussianCloudPipeline { ], }); + #[cfg(not(feature = "morph_particles"))] + let read_only = true; + + #[cfg(feature = "morph_particles")] + let read_only = false; + let gaussian_cloud_layout = render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { label: Some("gaussian_cloud_layout"), entries: &[ @@ -341,7 +347,7 @@ impl FromWorld for GaussianCloudPipeline { binding: 0, visibility: ShaderStages::all(), ty: BindingType::Buffer { - ty: BufferBindingType::Storage { read_only: false }, + ty: BufferBindingType::Storage { read_only }, has_dynamic_offset: false, min_binding_size: BufferSize::new(std::mem::size_of::() as u64), }, @@ -465,6 +471,9 @@ pub fn shader_defs( shader_defs.push("VISUALIZE_BOUNDING_BOX".into()); } + #[cfg(feature = "morph_particles")] + shader_defs.push("READ_WRITE_POINTS".into()); + shader_defs } diff --git a/src/sort/mod.rs b/src/sort/mod.rs index f24662b0..f26e1168 100644 --- a/src/sort/mod.rs +++ b/src/sort/mod.rs @@ -34,11 +34,15 @@ pub mod radix; #[cfg(feature = "sort_rayon")] pub mod rayon; +#[cfg(feature = "sort_std")] +pub mod std; + assert_cfg!( any( feature = "sort_radix", feature = "sort_rayon", + feature = "sort_std", ), "no sort mode enabled", ); @@ -59,6 +63,9 @@ pub enum SortMode { #[cfg(feature = "sort_rayon")] Rayon, + + #[cfg(feature = "sort_std")] + Std, } impl Default for SortMode { @@ -70,6 +77,9 @@ impl Default for SortMode { #[cfg(feature = "sort_radix")] return Self::Radix; + #[cfg(feature = "sort_std")] + return Self::Std; + Self::None } } @@ -86,6 +96,9 @@ impl Plugin for SortPlugin { #[cfg(feature = "sort_rayon")] app.add_plugins(rayon::RayonSortPlugin); + #[cfg(feature = "sort_std")] + app.add_plugins(std::StdSortPlugin); + app.register_type::(); app.init_asset::(); diff --git a/src/sort/std.rs b/src/sort/std.rs new file mode 100644 index 00000000..3bafa2fa --- /dev/null +++ b/src/sort/std.rs @@ -0,0 +1,102 @@ +use bevy::{ + prelude::*, + asset::LoadState, + utils::Instant, +}; + +use crate::{ + GaussianCloud, + GaussianCloudSettings, + sort::{ + SortedEntries, + SortMode, + }, +}; + + +#[derive(Default)] +pub struct StdSortPlugin; + +impl Plugin for StdSortPlugin { + fn build(&self, app: &mut App) { + app.add_systems(Update, std_sort); + } +} + +pub fn std_sort( + asset_server: Res, + gaussian_clouds_res: Res>, + mut sorted_entries_res: ResMut>, + gaussian_clouds: Query<( + &Handle, + &Handle, + &GaussianCloudSettings, + )>, + cameras: Query<( + &GlobalTransform, + &Camera3d, + )>, + mut last_camera_position: Local, + mut last_sort_time: Local>, +) { + let period = std::time::Duration::from_millis(100); + if let Some(last_sort_time) = last_sort_time.as_ref() { + if last_sort_time.elapsed() < period { + return; + } + } + + for ( + camera_transform, + _camera, + ) in cameras.iter() { + let camera_position = camera_transform.compute_transform().translation; + if *last_camera_position == camera_position { + return; + } + + for ( + gaussian_cloud_handle, + sorted_entries_handle, + settings, + ) in gaussian_clouds.iter() { + if settings.sort_mode != SortMode::Std { + continue; + } + + if Some(LoadState::Loading) == asset_server.get_load_state(gaussian_cloud_handle) { + continue; + } + + if Some(LoadState::Loading) == asset_server.get_load_state(sorted_entries_handle) { + continue; + } + + if let Some(gaussian_cloud) = gaussian_clouds_res.get(gaussian_cloud_handle) { + if let Some(sorted_entries) = sorted_entries_res.get_mut(sorted_entries_handle) { + assert_eq!(gaussian_cloud.gaussians.len(), sorted_entries.sorted.len()); + + *last_camera_position = camera_position; + *last_sort_time = Some(Instant::now()); + + gaussian_cloud.gaussians.iter() + .zip(sorted_entries.sorted.iter_mut()) + .enumerate() + .for_each(|(idx, (gaussian, sort_entry))| { + let position = Vec3::from_slice(gaussian.position.as_ref()); + let delta = camera_position - position; + + sort_entry.key = bytemuck::cast(delta.length_squared()); + sort_entry.index = idx as u32; + }); + + sorted_entries.sorted.sort_unstable_by(|a, b| { + bytemuck::cast::(b.key).partial_cmp(&bytemuck::cast::(a.key)).unwrap() + }); + + // TODO: update DrawIndirect buffer during sort phase (GPU sort will override default DrawIndirect) + } + } + } + } +} diff --git a/src/utils.rs b/src/utils.rs index e6781970..24b72910 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,5 +1,10 @@ +// #[cfg(target_arch = "wasm32")] +// pub use wasm_bindgen_rayon::init_thread_pool; + +#[cfg(target_arch = "wasm32")] +use wasm_bindgen::prelude::*; #[cfg(target_arch = "wasm32")] -pub use wasm_bindgen_rayon::init_thread_pool; +use std::collections::HashMap; pub fn setup_hooks() { @@ -9,3 +14,26 @@ pub fn setup_hooks() { console_error_panic_hook::set_once(); } } + + +#[cfg(not(target_arch = "wasm32"))] +pub fn get_arg(n: usize) -> Option { + std::env::args().nth(n) +} + +#[cfg(target_arch = "wasm32")] +pub fn get_arg(n: usize) -> Option { + let window = web_sys::window()?; // Get the window object + let location = window.location(); // Get the location object from the window + let search = location.search().ok()?; // Get the search (query string) part of the URL + + let args = search + .trim_start_matches('?') + .split('&') + .map(|s| s.splitn(2, '=').collect::>()) + .filter(|v| v.len() == 2) + .map(|v| (v[0].to_string(), v[1].to_string())) + .collect::>(); + + args.get(&format!("arg{}", n)).cloned() +} diff --git a/tests/gpu/gaussian.rs b/tests/gpu/gaussian.rs index 506e3299..03f0ba99 100644 --- a/tests/gpu/gaussian.rs +++ b/tests/gpu/gaussian.rs @@ -7,6 +7,7 @@ use bevy::{ prelude::*, app::AppExit, core::FrameCount, + core_pipeline::tonemapping::Tonemapping, render::view::screenshot::ScreenshotManager, window::PrimaryWindow, }; @@ -54,6 +55,7 @@ fn setup( commands.spawn(( Camera3dBundle { transform: Transform::from_translation(Vec3::new(0.0, 1.5, 5.0)), + tonemapping: Tonemapping::None, ..default() }, )); diff --git a/tests/gpu/radix.rs b/tests/gpu/radix.rs index d072b8ae..d448099f 100644 --- a/tests/gpu/radix.rs +++ b/tests/gpu/radix.rs @@ -9,9 +9,12 @@ use std::{ use bevy::{ prelude::*, core::FrameCount, - core_pipeline::core_3d::{ - CORE_3D, - Transparent3d, + core_pipeline::{ + core_3d::{ + CORE_3D, + Transparent3d, + }, + tonemapping::Tonemapping, }, render::{ RenderApp, @@ -97,6 +100,7 @@ fn setup( commands.spawn(( Camera3dBundle { transform: Transform::from_translation(Vec3::new(0.0, 1.5, 5.0)), + tonemapping: Tonemapping::None, ..default() }, )); diff --git a/tools/compare_aabb_obb.rs b/tools/compare_aabb_obb.rs index 4e963cfe..d48dea99 100644 --- a/tools/compare_aabb_obb.rs +++ b/tools/compare_aabb_obb.rs @@ -2,6 +2,7 @@ use bevy::{ prelude::*, app::AppExit, core::Name, + core_pipeline::tonemapping::Tonemapping, }; use bevy_inspector_egui::quick::WorldInspectorPlugin; use bevy_panorbit_camera::{ @@ -110,6 +111,7 @@ pub fn setup_aabb_obb_compare( commands.spawn(( Camera3dBundle { transform: Transform::from_translation(Vec3::new(0.0, 1.5, 5.0)), + tonemapping: Tonemapping::None, ..default() }, PanOrbitCamera{ diff --git a/viewer/viewer.rs b/viewer/viewer.rs index 2854ffe8..aa03ad84 100644 --- a/viewer/viewer.rs +++ b/viewer/viewer.rs @@ -2,6 +2,7 @@ use bevy::{ prelude::*, app::AppExit, core::Name, + core_pipeline::tonemapping::Tonemapping, diagnostic::{ DiagnosticsStore, FrameTimeDiagnosticsPlugin, @@ -17,12 +18,17 @@ use bevy_gaussian_splatting::{ GaussianCloud, GaussianSplattingBundle, GaussianSplattingPlugin, - morph::{ - ParticleBehaviors, - random_particle_behaviors, - }, random_gaussians, - utils::setup_hooks, + utils::{ + get_arg, + setup_hooks, + }, +}; + +#[cfg(feature = "morph_particles")] +use bevy_gaussian_splatting::morph::particle::{ + ParticleBehaviors, + random_particle_behaviors, }; @@ -53,23 +59,14 @@ fn setup_gaussian_cloud( mut commands: Commands, asset_server: Res, mut gaussian_assets: ResMut>, - mut particle_behavior_assets: ResMut>, ) { let cloud: Handle; - let mut particle_behaviors = None; - // TODO: add proper GaussianSplattingViewer argument parsing - let file_arg = std::env::args().nth(1); + let file_arg = get_arg(1); if let Some(n) = file_arg.clone().and_then(|s| s.parse::().ok()) { println!("generating {} gaussians", n); cloud = gaussian_assets.add(random_gaussians(n)); - - let behavior_arg = std::env::args().nth(2); - if let Some(k) = behavior_arg.clone().and_then(|s| s.parse::().ok()) { - println!("generating {} particle behaviors", k); - particle_behaviors = particle_behavior_assets.add(random_particle_behaviors(k)).into(); - } } else if let Some(filename) = file_arg { if filename == "--help".to_string() { println!("usage: cargo run -- [filename | n]"); @@ -82,7 +79,7 @@ fn setup_gaussian_cloud( cloud = gaussian_assets.add(GaussianCloud::test_model()); } - let mut entity = commands.spawn(( + commands.spawn(( GaussianSplattingBundle { cloud, ..default() @@ -90,13 +87,10 @@ fn setup_gaussian_cloud( Name::new("gaussian_cloud"), )); - if let Some(particle_behaviors) = particle_behaviors { - entity.insert(particle_behaviors); - } - commands.spawn(( Camera3dBundle { transform: Transform::from_translation(Vec3::new(0.0, 1.5, 5.0)), + tonemapping: Tonemapping::None, ..default() }, PanOrbitCamera{ @@ -107,25 +101,70 @@ fn setup_gaussian_cloud( } +#[cfg(feature = "morph_particles")] +fn setup_particle_behavior( + mut commands: Commands, + mut particle_behavior_assets: ResMut>, + gaussian_cloud: Query<( + Entity, + &Handle, + Without>, + )>, +) { + if gaussian_cloud.is_empty() { + return; + } + + let mut particle_behaviors = None; + + let file_arg = get_arg(1); + if let Some(_n) = file_arg.clone().and_then(|s| s.parse::().ok()) { + let behavior_arg = get_arg(2); + if let Some(k) = behavior_arg.clone().and_then(|s| s.parse::().ok()) { + println!("generating {} particle behaviors", k); + particle_behaviors = particle_behavior_assets.add(random_particle_behaviors(k)).into(); + } + } + + if let Some(particle_behaviors) = particle_behaviors { + commands.entity(gaussian_cloud.single().0) + .insert(particle_behaviors); + } +} + + fn example_app() { let config = GaussianSplattingViewer::default(); let mut app = App::new(); + #[cfg(target_arch = "wasm32")] + let primary_window = Some(Window { + fit_canvas_to_parent: true, + mode: bevy::window::WindowMode::Windowed, + present_mode: bevy::window::PresentMode::AutoVsync, + prevent_default_event_handling: true, + title: config.name.clone(), + ..default() + }); + + #[cfg(not(target_arch = "wasm32"))] + let primary_window = Some(Window { + fit_canvas_to_parent: true, + mode: bevy::window::WindowMode::Windowed, + present_mode: bevy::window::PresentMode::AutoVsync, + prevent_default_event_handling: false, + resolution: (config.width, config.height).into(), + title: config.name.clone(), + ..default() + }); + // setup for gaussian viewer app app.insert_resource(ClearColor(Color::rgb_u8(0, 0, 0))); app.add_plugins( DefaultPlugins .set(ImagePlugin::default_nearest()) .set(WindowPlugin { - primary_window: Some(Window { - fit_canvas_to_parent: false, - mode: bevy::window::WindowMode::Windowed, - present_mode: bevy::window::PresentMode::AutoVsync, - prevent_default_event_handling: false, - resolution: (config.width, config.height).into(), - title: config.name.clone(), - ..default() - }), + primary_window, ..default() }), ); @@ -152,6 +191,9 @@ fn example_app() { app.add_plugins(GaussianSplattingPlugin); app.add_systems(Startup, setup_gaussian_cloud); + #[cfg(feature = "morph_particles")] + app.add_systems(Update, setup_particle_behavior); + app.run(); } diff --git a/www/README.md b/www/README.md new file mode 100644 index 00000000..5e5d7d68 --- /dev/null +++ b/www/README.md @@ -0,0 +1,15 @@ +# bevy_gaussian_splatting for web + +## wasm support + +to build wasm run: + +```bash +cargo build --target wasm32-unknown-unknown --release --no-default-features --features "io_flexbuffers sort_std viewer" +``` + +to generate bindings: +> `wasm-bindgen --out-dir ./www/out/ --target web ./target/wasm32-unknown-unknown/release/viewer.wasm` + + +open a live server of `index.html` and append args: `?arg1=[n | *.gcloud]` diff --git a/www/index.html b/www/index.html new file mode 100644 index 00000000..4007d9e5 --- /dev/null +++ b/www/index.html @@ -0,0 +1,23 @@ + + + + bevy_gaussian_splatting + + + + +