Skip to content

Commit

Permalink
feat: 2dgs
Browse files Browse the repository at this point in the history
  • Loading branch information
@mosure
mosure committed Oct 15, 2024
1 parent a0e350c commit 2009dc6
Show file tree
Hide file tree
Showing 3 changed files with 127 additions and 44 deletions.
93 changes: 68 additions & 25 deletions src/render/gaussian.wgsl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -122,17 +123,33 @@ fn get_entry(index: u32) -> Entry {
struct GaussianVertexOutput {
@builtin(position) position: vec4<f32>,
@location(0) color: vec4<f32>,
@location(1) conic: vec3<f32>,
@location(2) uv: vec2<f32>,
@location(1) uv: vec2<f32>,
#ifdef GAUSSIAN_3D
@location(2) conic: vec3<f32>,
@location(3) major_minor: vec2<f32>,
#else ifdef GAUSSIAN_SURFEL
@location(2) local_to_pixel_u: vec3<f32>,
@location(3) local_to_pixel_v: vec3<f32>,
@location(4) local_to_pixel_w: vec3<f32>,
@location(5) mean_2d: vec2<f32>,
@location(6) radius: vec2<f32>,
#endif
};
#else
struct GaussianVertexOutput {
@builtin(position) position: vec4<f32>,
@location(0) @interpolate(flat) color: vec4<f32>,
@location(1) @interpolate(flat) conic: vec3<f32>,
@location(2) @interpolate(linear) uv: vec2<f32>,
@location(1) @interpolate(linear) uv: vec2<f32>,
#ifdef GAUSSIAN_3D
@location(2) @interpolate(flat) conic: vec3<f32>,
@location(3) @interpolate(linear) major_minor: vec2<f32>,
#else ifdef GAUSSIAN_SURFEL
@location(2) @interpolate(flat) local_to_pixel_u: vec3<f32>,
@location(3) @interpolate(flat) local_to_pixel_v: vec3<f32>,
@location(4) @interpolate(flat) local_to_pixel_w: vec3<f32>,
@location(5) @interpolate(flat) mean_2d: vec2<f32>,
@location(6) @interpolate(flat) radius: vec2<f32>,
#endif
};
#endif

Expand Down Expand Up @@ -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;
Expand All @@ -463,23 +468,62 @@ 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<f32>(
projected_position.xy + bb.xy,
projected_position.zw
projected_position.zw,
);

return output;
}

@fragment
fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4<f32> {
// TODO: surfel accumulation

#ifdef USE_AABB
#ifdef GAUSSIAN_SURFEL
let radius = input.radius;
let mean_2d = input.mean_2d;
let aspect = vec2<f32>(
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<f32>(
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;
Expand All @@ -499,7 +543,7 @@ fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4<f32> {
#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) ||
Expand All @@ -509,13 +553,12 @@ fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4<f32> {
}
#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<f32>(
input.color.rgb * final_alpha,
final_alpha,
input.color.rgb * alpha,
alpha,
);
}
75 changes: 57 additions & 18 deletions src/render/surfel.wgsl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,28 +79,38 @@
#endif


// TODO: analytic projection
struct Surfel {
local_to_pixel: mat3x3<f32>,
mean_2d: vec2<f32>,
extent: vec2<f32>,
};


fn get_bounding_box_cov2d(
cov2d: vec3<f32>,
extent: vec2<f32>,
direction: vec2<f32>,
cutoff: f32,
) -> vec4<f32> {
let fitler_size = 0.707106;

let extent = sqrt(max(
vec2<f32>(1.e-4, 1.e-4),
vec2<f32>(cov2d.x, cov2d.z),
if extent.x < 1.e-4 || extent.y < 1.e-4 {
return vec4<f32>(0.0);
}

let radius = sqrt(extent);
let max_radius = vec2<f32>(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<f32>(
vec2<f32>(radius) / view.viewport.zw,
max_radius / view.viewport.zw,
);

return vec4<f32>(
radius_ndc * direction,
radius * direction,
max_radius,
);
}

Expand All @@ -109,7 +119,9 @@ fn compute_cov2d_surfel(
gaussian_position: vec3<f32>,
index: u32,
cutoff: f32,
) -> vec3<f32> {
) -> Surfel {
var output: Surfel;

let rotation = get_rotation(index);
let scale = get_scale(index);

Expand All @@ -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<f32>(
vec4<f32>(L.x, 0.0),
Expand All @@ -137,22 +149,49 @@ fn compute_cov2d_surfel(

let test = vec3<f32>(cutoff * cutoff, cutoff * cutoff, -1.0);
let d = dot(test * T[2], T[2]);
if abs(d) < 1.0e-6 {
return vec3<f32>(0.0, 0.0, 0.0);
if abs(d) < 1.0e-4 {
output.extent = vec2<f32>(0.0);
return output;
}

let f = (1.0 / d) * test;
let means2d = vec2<f32>(
dot(f * T[0], T[2]),
dot(f * T[1], T[2]),
let mean_2d = vec2<f32>(
dot(f, T[0] * T[2]),
dot(f, T[1] * T[2]),
);

let t = vec2<f32>(
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<f32>,
pixel_coord: vec2<f32>,
mean_2d: vec2<f32>,
) -> 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<f32>(extent.x, covariance, extent.y);
return power;
}
3 changes: 2 additions & 1 deletion tools/surfel_plane.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit 2009dc6

Please sign in to comment.