diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index b438d6b7..68a54faf 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -18,6 +18,7 @@ #import bevy_gaussian_splatting::surfel::{ compute_cov2d_surfel, get_bounding_box_cov2d, + surfel_fragment_power, } #import bevy_gaussian_splatting::transform::{ world_to_clip, @@ -122,17 +123,33 @@ fn get_entry(index: u32) -> Entry { struct GaussianVertexOutput { @builtin(position) position: vec4, @location(0) color: vec4, - @location(1) conic: vec3, - @location(2) uv: vec2, + @location(1) uv: vec2, +#ifdef GAUSSIAN_3D + @location(2) conic: vec3, @location(3) major_minor: vec2, +#else ifdef GAUSSIAN_SURFEL + @location(2) local_to_pixel_u: vec3, + @location(3) local_to_pixel_v: vec3, + @location(4) local_to_pixel_w: vec3, + @location(5) mean_2d: vec2, + @location(6) radius: vec2, +#endif }; #else struct GaussianVertexOutput { @builtin(position) position: vec4, @location(0) @interpolate(flat) color: vec4, - @location(1) @interpolate(flat) conic: vec3, - @location(2) @interpolate(linear) uv: vec2, + @location(1) @interpolate(linear) uv: vec2, +#ifdef GAUSSIAN_3D + @location(2) @interpolate(flat) conic: vec3, @location(3) @interpolate(linear) major_minor: vec2, +#else ifdef GAUSSIAN_SURFEL + @location(2) @interpolate(flat) local_to_pixel_u: vec3, + @location(3) @interpolate(flat) local_to_pixel_v: vec3, + @location(4) @interpolate(flat) local_to_pixel_w: vec3, + @location(5) @interpolate(flat) mean_2d: vec2, + @location(6) @interpolate(flat) radius: vec2, +#endif }; #endif @@ -437,18 +454,6 @@ fn vs_points( quad_offset, cutoff, ); -#else ifdef GAUSSIAN_SURFEL - let cov2d = compute_cov2d_surfel( - transformed_position, - splat_index, - cutoff, - ); - let bb = get_bounding_box_cov2d( - cov2d, - quad_offset, - cutoff, - ); -#endif #ifdef USE_AABB let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y; @@ -463,10 +468,30 @@ fn vs_points( output.major_minor = bb.zw; #endif +#else ifdef GAUSSIAN_SURFEL + let surfel = compute_cov2d_surfel( + transformed_position, + splat_index, + cutoff, + ); + + output.local_to_pixel_u = surfel.local_to_pixel.x; + output.local_to_pixel_v = surfel.local_to_pixel.y; + output.local_to_pixel_w = surfel.local_to_pixel.z; + output.mean_2d = surfel.mean_2d; + + let bb = get_bounding_box_cov2d( + surfel.extent, + quad_offset, + cutoff, + ); + output.radius = bb.zw; +#endif + output.uv = quad_offset; output.position = vec4( projected_position.xy + bb.xy, - projected_position.zw + projected_position.zw, ); return output; @@ -474,12 +499,31 @@ fn vs_points( @fragment fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { - // TODO: surfel accumulation - #ifdef USE_AABB +#ifdef GAUSSIAN_SURFEL + let radius = input.radius; + let mean_2d = input.mean_2d; + let aspect = vec2( + 1.0, + view.viewport.z / view.viewport.w, + ); + let pixel_coord = input.uv * radius * aspect + mean_2d; + // let pixel_coord = input.position.xy * view.viewport.zw + view.viewport.xy; + + let power = surfel_fragment_power( + mat3x3( + input.local_to_pixel_u, + input.local_to_pixel_v, + input.local_to_pixel_w, + ), + pixel_coord, + mean_2d, + ); +#else ifdef GAUSSIAN_3D 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; +#endif if (power > 0.0) { discard; @@ -499,7 +543,7 @@ fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { #endif #ifdef VISUALIZE_BOUNDING_BOX - let uv = (input.uv + 1.0) / 2.0; + let uv = input.uv * 0.5 + 0.5; let edge_width = 0.08; if ( (uv.x < edge_width || uv.x > 1.0 - edge_width) || @@ -509,13 +553,12 @@ fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { } #endif - let alpha = exp(power); - let final_alpha = alpha * input.color.a; + let alpha = min(exp(power) * input.color.a, 0.999); - // TODO: round final_alpha to terminate depth test? + // TODO: round alpha to terminate depth test? return vec4( - input.color.rgb * final_alpha, - final_alpha, + input.color.rgb * alpha, + alpha, ); } diff --git a/src/render/surfel.wgsl b/src/render/surfel.wgsl index d97bfd17..7d9ab1e5 100644 --- a/src/render/surfel.wgsl +++ b/src/render/surfel.wgsl @@ -79,28 +79,38 @@ #endif -// TODO: analytic projection +struct Surfel { + local_to_pixel: mat3x3, + mean_2d: vec2, + extent: vec2, +}; + + fn get_bounding_box_cov2d( - cov2d: vec3, + extent: vec2, direction: vec2, cutoff: f32, ) -> vec4 { let fitler_size = 0.707106; - let extent = sqrt(max( - vec2(1.e-4, 1.e-4), - vec2(cov2d.x, cov2d.z), + if extent.x < 1.e-4 || extent.y < 1.e-4 { + return vec4(0.0); + } + + let radius = sqrt(extent); + let max_radius = vec2(max( + max(radius.x, radius.y), + cutoff * fitler_size, )); - let radius = ceil(max(max(extent.x, extent.y), cutoff * fitler_size)); // TODO: verify OBB capability let radius_ndc = vec2( - vec2(radius) / view.viewport.zw, + max_radius / view.viewport.zw, ); return vec4( radius_ndc * direction, - radius * direction, + max_radius, ); } @@ -109,7 +119,9 @@ fn compute_cov2d_surfel( gaussian_position: vec3, index: u32, cutoff: f32, -) -> vec3 { +) -> Surfel { + var output: Surfel; + let rotation = get_rotation(index); let scale = get_scale(index); @@ -122,7 +134,7 @@ fn compute_cov2d_surfel( let S = get_scale_matrix(scale); let R = get_rotation_matrix(rotation); - let L = R * S;// * transpose(T_r); + let L = T_r * S * R; let world_from_local = mat3x4( vec4(L.x, 0.0), @@ -137,22 +149,49 @@ fn compute_cov2d_surfel( let test = vec3(cutoff * cutoff, cutoff * cutoff, -1.0); let d = dot(test * T[2], T[2]); - if abs(d) < 1.0e-6 { - return vec3(0.0, 0.0, 0.0); + if abs(d) < 1.0e-4 { + output.extent = vec2(0.0); + return output; } let f = (1.0 / d) * test; - let means2d = vec2( - dot(f * T[0], T[2]), - dot(f * T[1], T[2]), + let mean_2d = vec2( + dot(f, T[0] * T[2]), + dot(f, T[1] * T[2]), ); let t = vec2( dot(f * T[0], T[0]), dot(f * T[1], T[1]), ); - let extent = means2d * means2d - t; - let covariance = means2d.x * means2d.y - dot(f * T[0], T[1]); + let extent = mean_2d * mean_2d - t; + + output.local_to_pixel = T; + output.mean_2d = mean_2d; + output.extent = extent; + return output; +} + +fn surfel_fragment_power( + local_to_pixel: mat3x3, + pixel_coord: vec2, + mean_2d: vec2, +) -> f32 { + let deltas = mean_2d - pixel_coord; + + let hu = pixel_coord.x * local_to_pixel.z - local_to_pixel.x; + let hv = pixel_coord.y * local_to_pixel.z - local_to_pixel.y; + + let p = cross(hu, hv); + + let us = p.x / p.z; + let vs = p.y / p.z; + + let sigmas_3d = us * us + vs * vs; + let sigmas_2d = 2.0 * (deltas.x * deltas.x + deltas.y * deltas.y); + + let sigmas = 0.5 * min(sigmas_3d, sigmas_2d); + let power = -sigmas; - return vec3(extent.x, covariance, extent.y); + return power; } diff --git a/tools/surfel_plane.rs b/tools/surfel_plane.rs index 033f4766..9d54bb5e 100644 --- a/tools/surfel_plane.rs +++ b/tools/surfel_plane.rs @@ -47,7 +47,7 @@ pub fn setup_surfel_compare( let x = i as f32 * spacing - (grid_size_x as f32 * spacing) / 2.0; let y = j as f32 * spacing - (grid_size_y as f32 * spacing) / 2.0; let position = [x, y, 0.0, 1.0]; - let scale = [1.0, 1.0, 0.1, 0.5]; + let scale = [1.0, 1.0, 0.01, 0.5]; let gaussian = Gaussian { position_visibility: position.into(), @@ -100,6 +100,7 @@ pub fn setup_surfel_compare( cloud: gaussian_assets.add(GaussianCloud::from_gaussians(red_gaussians)), settings: GaussianCloudSettings { visualize_bounding_box, + aabb: true, transform: Transform::from_translation(Vec3::new(spacing, spacing, 0.0)), gaussian_mode: GaussianMode::GaussianSurfel, ..default()