diff --git a/Cargo.lock b/Cargo.lock index 0b6052a0..f2c51889 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1246,7 +1246,7 @@ dependencies = [ [[package]] name = "bevy_gaussian_splatting" -version = "4.1.1" +version = "4.2.0" dependencies = [ "base64 0.22.1", "bevy 0.15.2", @@ -1317,6 +1317,7 @@ dependencies = [ "bevy_gizmos_macros 0.15.2", "bevy_image", "bevy_math 0.15.2", + "bevy_pbr 0.15.2", "bevy_reflect 0.15.2", "bevy_render 0.15.2", "bevy_sprite", diff --git a/Cargo.toml b/Cargo.toml index 06b8faef..9aba06b2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "bevy_gaussian_splatting" description = "bevy gaussian splatting render pipeline plugin" -version = "4.1.1" +version = "4.2.0" edition = "2021" authors = ["mosure "] license = "MIT OR Apache-2.0" @@ -186,6 +186,7 @@ default-features = false features = [ "bevy_asset", "bevy_core_pipeline", + "bevy_pbr", "bevy_render", "bevy_winit", "x11", diff --git a/src/gaussian/settings.rs b/src/gaussian/settings.rs index 1181a147..41db6326 100644 --- a/src/gaussian/settings.rs +++ b/src/gaussian/settings.rs @@ -86,7 +86,7 @@ pub enum RasterizeMode { #[default] Color, Depth, - Flow, + OpticalFlow, Normal, Velocity, } diff --git a/src/io/ply.rs b/src/io/ply.rs index 9d80976a..97b3abe8 100644 --- a/src/io/ply.rs +++ b/src/io/ply.rs @@ -98,7 +98,7 @@ pub fn parse_ply_3d( let required_properties = vec![ "x", "y", "z", "f_dc_0", "f_dc_1", "f_dc_2", - "scale_0", "scale_1", "scale_2", + "scale_0", "scale_1", "opacity", "rot_0", "rot_1", "rot_2", "rot_3", ]; diff --git a/src/material/mod.rs b/src/material/mod.rs index b940f137..2d1cf194 100644 --- a/src/material/mod.rs +++ b/src/material/mod.rs @@ -2,6 +2,7 @@ use bevy::prelude::*; pub mod classification; pub mod depth; +pub mod optical_flow; pub mod spherical_harmonics; pub mod spherindrical_harmonics; @@ -21,6 +22,7 @@ impl Plugin for MaterialPlugin { app.add_plugins(( classification::ClassificationMaterialPlugin, depth::DepthMaterialPlugin, + optical_flow::OpticalFlowMaterialPlugin, spherical_harmonics::SphericalHarmonicCoefficientsPlugin, spherindrical_harmonics::SpherindricalHarmonicCoefficientsPlugin, )); diff --git a/src/material/optical_flow.rs b/src/material/optical_flow.rs new file mode 100644 index 00000000..e50d4621 --- /dev/null +++ b/src/material/optical_flow.rs @@ -0,0 +1,21 @@ +use bevy::{ + prelude::*, + asset::load_internal_asset, +}; + + +const OPTICAL_FLOW_SHADER_HANDLE: Handle = Handle::weak_from_u128(1451151234); + + +pub struct OpticalFlowMaterialPlugin; + +impl Plugin for OpticalFlowMaterialPlugin { + fn build(&self, app: &mut App) { + load_internal_asset!( + app, + OPTICAL_FLOW_SHADER_HANDLE, + "optical_flow.wgsl", + Shader::from_wgsl + ); + } +} diff --git a/src/material/optical_flow.wgsl b/src/material/optical_flow.wgsl new file mode 100644 index 00000000..83249418 --- /dev/null +++ b/src/material/optical_flow.wgsl @@ -0,0 +1,59 @@ +#define_import_path bevy_gaussian_splatting::optical_flow + +#import bevy_pbr::{ + forward_io::VertexOutput, + prepass_utils, +} +#import bevy_render::color_operations::hsv_to_rgb +#import bevy_render::maths::PI_2 + +#import bevy_gaussian_splatting::bindings::{ + globals, + previous_view_uniforms, + view, +} + + + +fn calculate_motion_vector( + world_position: vec3, + previous_world_position: vec3, +) -> vec2 { + let world_position_t = vec4(world_position, 1.0); + let previous_world_position_t = vec4(previous_world_position, 1.0); + let clip_position_t = view.unjittered_clip_from_world * world_position_t; + let clip_position = clip_position_t.xy / clip_position_t.w; + let previous_clip_position_t = previous_view_uniforms.clip_from_world * previous_world_position_t; + let previous_clip_position = previous_clip_position_t.xy / previous_clip_position_t.w; + // These motion vectors are used as offsets to UV positions and are stored + // in the range -1,1 to allow offsetting from the one corner to the + // diagonally-opposite corner in UV coordinates, in either direction. + // A difference between diagonally-opposite corners of clip space is in the + // range -2,2, so this needs to be scaled by 0.5. And the V direction goes + // down where clip space y goes up, so y needs to be flipped. + return (clip_position - previous_clip_position) * vec2(0.5, -0.5); +} + + +fn optical_flow_to_rgb( + motion_vector: vec2, +) -> vec3 { + let flow = motion_vector / globals.delta_time; + + let radius = length(flow); + var angle = atan2(flow.y, flow.x); + if (angle < 0.0) { + angle += PI_2; + } + + // let sigma: f32 = 0.15; + // let norm_factor = sigma * 2.0; + // let m = clamp(radius / norm_factor, 0.0, 1.0); + let m = clamp(radius, 0.0, 1.0); + + let rgb = hsv_to_rgb(vec3(angle, m, 1.0)); + return rgb; +} + +// TODO: support immediate vs. persistent previous_view, aiding with no-smoothness on the pan-orbit camera (required by cpu sort) +// TODO: set clear color to white in optical flow render mode diff --git a/src/morph/particle.rs b/src/morph/particle.rs index 21179217..648794fc 100644 --- a/src/morph/particle.rs +++ b/src/morph/particle.rs @@ -261,7 +261,7 @@ impl FromWorld for ParticleBehaviorPipeline { 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.compute_view_layout.clone(), gaussian_cloud_pipeline.gaussian_uniform_layout.clone(), gaussian_cloud_pipeline.gaussian_cloud_layout.clone(), particle_behavior_layout.clone(), diff --git a/src/render/bindings.wgsl b/src/render/bindings.wgsl index 2ff89cf8..6401c609 100644 --- a/src/render/bindings.wgsl +++ b/src/render/bindings.wgsl @@ -1,11 +1,13 @@ #define_import_path bevy_gaussian_splatting::bindings +#import bevy_pbr::prepass_bindings::PreviousViewUniforms #import bevy_render::globals::Globals #import bevy_render::view::View @group(0) @binding(0) var view: View; @group(0) @binding(1) var globals: Globals; +@group(0) @binding(2) var previous_view_uniforms: PreviousViewUniforms; struct GaussianUniforms { transform: mat4x4, diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index 08e54e18..6c5b7802 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -5,6 +5,10 @@ } #import bevy_gaussian_splatting::classification::class_to_rgb #import bevy_gaussian_splatting::depth::depth_to_rgb +#import bevy_gaussian_splatting::optical_flow::{ + calculate_motion_vector, + optical_flow_to_rgb, +} #import bevy_gaussian_splatting::helpers::{ get_rotation_matrix, get_scale_matrix, @@ -178,6 +182,7 @@ fn vs_points( let position = vec4(get_position(splat_index), 1.0); var transformed_position = (gaussian_uniforms.transform * position).xyz; + var previous_transformed_position = transformed_position; #ifdef DRAW_SELECTED discard_quad |= get_visibility(splat_index) < 0.5; @@ -254,6 +259,7 @@ fn vs_points( let position_t = vec4(position.xyz + gaussian_4d.delta_mean, 1.0); transformed_position = (gaussian_uniforms.transform * position_t).xyz; + // TODO: set previous_transformed_position based on temporal position delta let projected_position = world_to_clip(transformed_position); if !in_frustum(projected_position.xyz) { @@ -346,9 +352,13 @@ fn vs_points( 0.5 * (t.y + 1.0), 0.5 * (t.z + 1.0) ); -#else ifdef RASTERIZE_FLOW - // TODO: optical flow rendering - rgb = vec3(1.0); +#else ifdef RASTERIZE_OPTICAL_FLOW + let motion_vector = calculate_motion_vector( + transformed_position, + previous_transformed_position, + ); + + rgb = optical_flow_to_rgb(motion_vector); #else ifdef RASTERIZE_VELOCITY let time_delta = 1e-3; let future_gaussian_4d = conditional_cov3d( diff --git a/src/render/mod.rs b/src/render/mod.rs index 8e19d609..250fe23d 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -3,7 +3,15 @@ use std::hash::Hash; use bevy::{ prelude::*, asset::load_internal_asset, - core_pipeline::core_3d::Transparent3d, + core_pipeline::{ + core_3d::Transparent3d, + prepass::{ + MotionVectorPrepass, + PreviousViewData, + PreviousViewUniforms, + PreviousViewUniformOffset, + }, + }, ecs::{ query::ROQueryItem, system::{ @@ -11,6 +19,7 @@ use bevy::{ SystemParamItem, } }, + pbr::PrepassViewBindGroup, render::{ Extract, extract_component::{ @@ -381,6 +390,7 @@ pub struct CloudPipeline { pub gaussian_cloud_layout: BindGroupLayout, pub gaussian_uniform_layout: BindGroupLayout, pub view_layout: BindGroupLayout, + pub compute_view_layout: BindGroupLayout, pub sorted_layout: BindGroupLayout, phantom: std::marker::PhantomData, } @@ -389,7 +399,41 @@ impl FromWorld for CloudPipeline { fn from_world(render_world: &mut World) -> Self { let render_device = render_world.resource::(); + // TODO: store both ShaderStages::all() and ShaderStages::VERTEX_FRAGMENT pipelines (for previous_view/sort nodes) let view_layout_entries = vec![ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::VERTEX_FRAGMENT, + ty: BindingType::Buffer { + ty: BufferBindingType::Uniform, + has_dynamic_offset: true, + min_binding_size: Some(ViewUniform::min_size()), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::VERTEX_FRAGMENT, + ty: BindingType::Buffer { + ty: BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: Some(GlobalsUniform::min_size()), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::VERTEX_FRAGMENT, + ty: BindingType::Buffer { + ty: BufferBindingType::Uniform, + has_dynamic_offset: true, + min_binding_size: Some(PreviousViewData::min_size()), + }, + count: None, + }, + ]; + + let compute_view_layout_entries = vec![ BindGroupLayoutEntry { binding: 0, visibility: ShaderStages::all(), @@ -410,6 +454,16 @@ impl FromWorld for CloudPipeline { }, count: None, }, + // BindGroupLayoutEntry { + // binding: 2, + // visibility: ShaderStages::all(), + // ty: BindingType::Buffer { + // ty: BufferBindingType::Uniform, + // has_dynamic_offset: true, + // min_binding_size: Some(PreviousViewData::min_size()), + // }, + // count: None, + // }, ]; let view_layout = render_device.create_bind_group_layout( @@ -417,6 +471,11 @@ impl FromWorld for CloudPipeline { &view_layout_entries, ); + let compute_view_layout = render_device.create_bind_group_layout( + Some("gaussian_compute_view_layout"), + &compute_view_layout_entries, + ); + let gaussian_uniform_layout = render_device.create_bind_group_layout( Some("gaussian_uniform_layout"), &[ @@ -465,6 +524,7 @@ impl FromWorld for CloudPipeline { gaussian_cloud_layout, gaussian_uniform_layout, view_layout, + compute_view_layout, shader: GAUSSIAN_SHADER_HANDLE, sorted_layout, phantom: std::marker::PhantomData, @@ -624,7 +684,7 @@ pub fn shader_defs( RasterizeMode::Classification => shader_defs.push("RASTERIZE_CLASSIFICATION".into()), RasterizeMode::Color => shader_defs.push("RASTERIZE_COLOR".into()), RasterizeMode::Depth => shader_defs.push("RASTERIZE_DEPTH".into()), - RasterizeMode::Flow => shader_defs.push("RASTERIZE_FLOW".into()), + RasterizeMode::OpticalFlow => shader_defs.push("RASTERIZE_OPTICAL_FLOW".into()), RasterizeMode::Normal => shader_defs.push("RASTERIZE_NORMAL".into()), RasterizeMode::Velocity => shader_defs.push("RASTERIZE_VELOCITY".into()), } @@ -726,7 +786,8 @@ impl SpecializedRenderPipeline for CloudPipeline { type DrawGaussians = ( SetItemPipeline, - SetGaussianViewBindGroup<0>, + // SetViewBindGroup<0>, + SetPreviousViewBindGroup<0>, SetGaussianUniformBindGroup<1>, DrawGaussianInstanced, ); @@ -949,10 +1010,12 @@ pub fn queue_gaussian_view_bind_groups( render_device: Res, gaussian_cloud_pipeline: Res>, view_uniforms: Res, + previous_view_uniforms: Res, views: Query< ( Entity, &ExtractedView, + Option<&PreviousViewData>, ), With, >, @@ -960,14 +1023,17 @@ pub fn queue_gaussian_view_bind_groups( ) { if let ( Some(view_binding), + Some(previous_view_binding), Some(globals), ) = ( view_uniforms.uniforms.binding(), + previous_view_uniforms.uniforms.binding(), globals_buffer.buffer.binding(), ) { for ( entity, _extracted_view, + _maybe_previous_view, ) in &views { let layout = &gaussian_cloud_pipeline.view_layout; @@ -980,6 +1046,10 @@ pub fn queue_gaussian_view_bind_groups( binding: 1, resource: globals.clone(), }, + BindGroupEntry { + binding: 2, + resource: previous_view_binding.clone(), + }, ]; let view_bind_group = render_device.create_bind_group( @@ -990,15 +1060,17 @@ pub fn queue_gaussian_view_bind_groups( debug!("inserting gaussian view bind group"); - commands.entity(entity).insert(GaussianViewBindGroup { - value: view_bind_group, - }); + commands + .entity(entity) + .insert(GaussianViewBindGroup { + value: view_bind_group, + }); } } } -pub struct SetGaussianViewBindGroup; -impl RenderCommand

for SetGaussianViewBindGroup { +pub struct SetViewBindGroup; +impl RenderCommand

for SetViewBindGroup { type Param = (); type ViewQuery = ( Read, @@ -1008,7 +1080,7 @@ impl RenderCommand

for SetGaussianViewBindGroup #[inline] fn render<'w>( - _item: &P, + _: &P, ( gaussian_view_bind_group, view_uniform, @@ -1026,7 +1098,54 @@ impl RenderCommand

for SetGaussianViewBindGroup &[view_uniform.offset], ); - debug!("set gaussian view bind group"); + debug!("set view bind group"); + + RenderCommandResult::Success + } +} + + +pub struct SetPreviousViewBindGroup; +impl RenderCommand

for SetPreviousViewBindGroup { + type Param = SRes; + type ViewQuery = ( + Read, + Has, + Option>, + ); + type ItemQuery = (); + + #[inline] + fn render<'w>( + _: &P, + ( + view_uniform_offset, + has_motion_vector_prepass, + previous_view_uniform_offset, + ): ROQueryItem< + 'w, + Self::ViewQuery, + >, + _entity: Option<()>, + prepass_view_bind_group: SystemParamItem<'w, '_, Self::Param>, + pass: &mut TrackedRenderPass<'w>, + ) -> RenderCommandResult { + let prepass_view_bind_group = prepass_view_bind_group.into_inner(); + match previous_view_uniform_offset { + Some(previous_view_uniform_offset) if has_motion_vector_prepass => { + pass.set_bind_group( + I, + prepass_view_bind_group.motion_vectors.as_ref().unwrap(), + &[ + view_uniform_offset.offset, + previous_view_uniform_offset.offset, + ], + ); + } + _ => {} + } + + debug!("set previous view bind group"); RenderCommandResult::Success } diff --git a/src/sort/mod.rs b/src/sort/mod.rs index 81869fb3..d32bb041 100644 --- a/src/sort/mod.rs +++ b/src/sort/mod.rs @@ -115,7 +115,7 @@ pub struct SortConfig { impl Default for SortConfig { fn default() -> Self { Self { - period_ms: 100, + period_ms: 1000, } } } @@ -244,6 +244,7 @@ fn update_sort_trigger( if camera_movement { sort_trigger.needs_sort = true; + sort_trigger.last_sort_time = Some(Instant::now()); sort_trigger.last_camera_position = camera_position; } } diff --git a/src/sort/radix.rs b/src/sort/radix.rs index 4f32af19..49760ed9 100644 --- a/src/sort/radix.rs +++ b/src/sort/radix.rs @@ -329,7 +329,7 @@ impl FromWorld for RadixSortPipeline { ); let sorting_layout = vec![ - gaussian_cloud_pipeline.view_layout.clone(), + gaussian_cloud_pipeline.compute_view_layout.clone(), gaussian_cloud_pipeline.gaussian_uniform_layout.clone(), gaussian_cloud_pipeline.gaussian_cloud_layout.clone(), radix_sort_layout.clone(), diff --git a/viewer/viewer.rs b/viewer/viewer.rs index fd9a4f09..2e4dfcd2 100644 --- a/viewer/viewer.rs +++ b/viewer/viewer.rs @@ -6,7 +6,10 @@ use bevy::{ app::AppExit, color::palettes::css::GOLD, core::FrameCount, - core_pipeline::tonemapping::Tonemapping, + core_pipeline::{ + prepass::MotionVectorPrepass, + tonemapping::Tonemapping, + }, diagnostic::{ DiagnosticsStore, FrameTimeDiagnosticsPlugin, @@ -76,7 +79,7 @@ fn parse_input_file( String::from_utf8(data).unwrap() }, Err(e) => { - warn!("failed to decode base64 input: {:?}", e); + debug!("failed to decode base64 input: {:?}", e); input_file.to_string() } }; @@ -151,11 +154,12 @@ fn setup_gaussian_cloud( Camera3d::default(), Transform::from_translation(Vec3::new(0.0, 1.5, 5.0)), Tonemapping::None, + MotionVectorPrepass, PanOrbitCamera { allow_upside_down: true, - orbit_smoothness: 0.0, - pan_smoothness: 0.0, - zoom_smoothness: 0.0, + orbit_smoothness: 0.1, + pan_smoothness: 0.1, + zoom_smoothness: 0.1, ..default() }, GaussianCamera::default(),