Bevy Shader Explorer
Search across packages, versions, files, declarations, and comments. Use multiple terms to narrow results (AND), and add @vertex, @fragment, or @compute to filter by shader stage.
Packages › bevy_gaussian_splatting › gaussian.wgsl › source
← Back to declarations 
1#import bevy_gaussian_splatting::bindings::{
2    view,
3    gaussian_uniforms,
4    sorting_pass_index,
5    sorting,
6    draw_indirect,
7    input_entries,
8    output_entries,
9    Entry,
11#import bevy_gaussian_splatting::depth::{
12    depth_to_rgb,
14#import bevy_gaussian_splatting::helpers::{
15    get_rotation_matrix,
16    get_scale_matrix,
18#import bevy_gaussian_splatting::transform::{
19    world_to_clip,
20    in_frustum,
23#ifdef GAUSSIAN_SURFEL
24#import bevy_gaussian_splatting::surfel::{
25    compute_cov2d_surfel,
26    get_bounding_box_cov2d,
27    surfel_fragment_power,
29#endif
31#ifdef PACKED
32#ifdef PRECOMPUTE_COVARIANCE_3D
33#import bevy_gaussian_splatting::packed::{
34    get_position,
35    get_color,
36    get_visibility,
37    get_opacity,
38    get_cov3d,
40#else
41#import bevy_gaussian_splatting::packed::{
42    get_position,
43    get_color,
44    get_visibility,
45    get_opacity,
46    get_rotation,
47    get_scale,
49#endif
50#else
52#ifdef BUFFER_STORAGE
53#ifdef PRECOMPUTE_COVARIANCE_3D
54#import bevy_gaussian_splatting::planar::{
55    get_position,
56    get_color,
57    get_visibility,
58    get_opacity,
59    get_cov3d,
61#else
62#import bevy_gaussian_splatting::planar::{
63    get_position,
64    get_color,
65    get_visibility,
66    get_opacity,
67    get_rotation,
68    get_scale,
70#endif
71#endif
73#endif
76#ifdef BUFFER_TEXTURE
77#ifdef PRECOMPUTE_COVARIANCE_3D
78#import bevy_gaussian_splatting::texture::{
79    get_position,
80    get_color,
81    get_visibility,
82    get_opacity,
83    get_cov3d,
84    location,
86#else
87#import bevy_gaussian_splatting::texture::{
88    get_position,
89    get_color,
90    get_visibility,
91    get_opacity,
92    get_rotation,
93    get_scale,
94    location,
96#endif
97#endif
100#ifdef BUFFER_STORAGE
101@group(3) @binding(0) var<storage, read> sorted_entries: array<Entry>;
103fn get_entry(index: u32) -> Entry {
104    return sorted_entries[index];
106#endif
108#ifdef BUFFER_TEXTURE
109@group(3) @binding(0) var sorted_entries: texture_2d<u32>;
111fn get_entry(index: u32) -> Entry {
112    let sample = textureLoad(
113        sorted_entries,
114        location(index),
115        0,
116    );
118    return Entry(
119        sample.r,
120        sample.g,
121    );
123#endif
125#ifdef WEBGL2
126struct GaussianVertexOutput {
127    @builtin(position) position: vec4<f32>,
128    @location(0) color: vec4<f32>,
129    @location(1) uv: vec2<f32>,
130#ifdef GAUSSIAN_3D
131    @location(2) conic: vec3<f32>,
132    @location(3) major_minor: vec2<f32>,
133#else ifdef GAUSSIAN_SURFEL
134    @location(2) local_to_pixel_u: vec3<f32>,
135    @location(3) local_to_pixel_v: vec3<f32>,
136    @location(4) local_to_pixel_w: vec3<f32>,
137    @location(5) mean_2d: vec2<f32>,
138    @location(6) radius: vec2<f32>,
139#endif
140};
141#else
142struct GaussianVertexOutput {
143    @builtin(position) position: vec4<f32>,
144    @location(0) @interpolate(flat) color: vec4<f32>,
145    @location(1) @interpolate(linear) uv: vec2<f32>,
146#ifdef GAUSSIAN_3D
147    @location(2) @interpolate(flat) conic: vec3<f32>,
148    @location(3) @interpolate(linear) major_minor: vec2<f32>,
149#else ifdef GAUSSIAN_SURFEL
150    @location(2) @interpolate(flat) local_to_pixel_u: vec3<f32>,
151    @location(3) @interpolate(flat) local_to_pixel_v: vec3<f32>,
152    @location(4) @interpolate(flat) local_to_pixel_w: vec3<f32>,
153    @location(5) @interpolate(flat) mean_2d: vec2<f32>,
154    @location(6) @interpolate(flat) radius: vec2<f32>,
155#endif
156};
157#endif
160// https://github.com/cvlab-epfl/gaussian-splatting-web/blob/905b3c0fb8961e42c79ef97e64609e82383ca1c2/src/shaders.ts#L185
161// TODO: precompute
162fn compute_cov3d(scale: vec3<f32>, rotation: vec4<f32>) -> array<f32, 6> {
163    let S = get_scale_matrix(scale);
165    let T = mat3x3<f32>(
166        gaussian_uniforms.transform[0].xyz,
167        gaussian_uniforms.transform[1].xyz,
168        gaussian_uniforms.transform[2].xyz,
169    );
171    let R = get_rotation_matrix(rotation);
173    let M = S * R;
174    let Sigma = transpose(M) * M;
175    let TS = T * Sigma * transpose(T);
177    return array<f32, 6>(
178        TS[0][0],
179        TS[0][1],
180        TS[0][2],
181        TS[1][1],
182        TS[1][2],
183        TS[2][2],
184    );
187fn compute_cov2d_3dgs(
188    position: vec3<f32>,
189    index: u32,
190) -> vec3<f32> {
191#ifdef PRECOMPUTE_COVARIANCE_3D
192    let cov3d = get_cov3d(index);
193#else
194    let rotation = get_rotation(index);
195    let scale = get_scale(index);
197    let cov3d = compute_cov3d(scale, rotation);
198#endif
200    let Vrk = mat3x3(
201        cov3d[0], cov3d[1], cov3d[2],
202        cov3d[1], cov3d[3], cov3d[4],
203        cov3d[2], cov3d[4], cov3d[5],
204    );
206    var t = view.view_from_world * vec4<f32>(position, 1.0);
208    let focal = vec2<f32>(
209        view.clip_from_view.x.x * view.viewport.z,
210        view.clip_from_view.y.y * view.viewport.w,
211    );
213    let s = 1.0 / (t.z * t.z);
214    let J = mat3x3(
215        focal.x / t.z, 0.0, -(focal.x * t.x) * s,
216        0.0, -focal.y / t.z, (focal.y * t.y) * s,
217        0.0, 0.0, 0.0,
218    );
220    let W = transpose(
221        mat3x3<f32>(
222            view.view_from_world.x.xyz,
223            view.view_from_world.y.xyz,
224            view.view_from_world.z.xyz,
225        )
226    );
228    let T = W * J;
230    var cov = transpose(T) * transpose(Vrk) * T;
231    cov[0][0] += 0.3f;
232    cov[1][1] += 0.3f;
234    return vec3<f32>(cov[0][0], cov[0][1], cov[1][1]);
237fn get_bounding_box(
238    cov2d: vec3<f32>,
239    direction: vec2<f32>,
240    cutoff: f32,
241) -> vec4<f32> {
242    // return vec4<f32>(offset, uv);
244    let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y;
245    let trace = cov2d.x + cov2d.z;
246    let mid = 0.5 * trace;
247    let discriminant = max(0.0, mid * mid - det);
249    let term = sqrt(discriminant);
251    let lambda1 = mid + term;
252    let lambda2 = max(mid - term, 0.0);
254    let x_axis_length = sqrt(lambda1);
255    let y_axis_length = sqrt(lambda2);
258#ifdef USE_AABB
259    let radius_px = cutoff * max(x_axis_length, y_axis_length);
260    let radius_ndc = vec2<f32>(
261        radius_px / view.viewport.zw,
262    );
264    return vec4<f32>(
265        radius_ndc * direction,
266        radius_px * direction,
267    );
268#endif
270#ifdef USE_OBB
272    let a = (cov2d.x - cov2d.z) * (cov2d.x - cov2d.z);
273    let b = sqrt(a + 4.0 * cov2d.y * cov2d.y);
274    let major_radius = sqrt((cov2d.x + cov2d.z + b) * 0.5);
275    let minor_radius = sqrt((cov2d.x + cov2d.z - b) * 0.5);
277    let bounds = cutoff * vec2<f32>(
278        major_radius,
279        minor_radius,
280    );
282    let eigvec1 = normalize(vec2<f32>(
283        -cov2d.y,
284        lambda1 - cov2d.x,
285    ));
286    let eigvec2 = vec2<f32>(
287        eigvec1.y,
288        -eigvec1.x
289    );
291    let rotation_matrix = transpose(
292        mat2x2(
293            eigvec1,
294            eigvec2,
295        )
296    );
298    let scaled_vertex = direction * bounds;
299    let rotated_vertex = scaled_vertex * rotation_matrix;
301    let scaling_factor = 1.0 / view.viewport.zw;
302    let ndc_vertex = rotated_vertex * scaling_factor;
304    return vec4<f32>(
305        ndc_vertex,
306        rotated_vertex,
307    );
308#endif
312// @compute @workgroup_size(#{RADIX_BASE}, #{RADIX_DIGIT_PLACES})
313// fn gaussian_compute(
314//     @builtin(local_invocation_id) gl_LocalInvocationID: vec3<u32>,
315//     @builtin(global_invocation_id) gl_GlobalInvocationID: vec3<u32>,
316// ) {
317//     // TODO: compute cov2d, color (any non-quad gaussian property)
318// }
320fn inverted_infinity_norm(v: vec3<f32>) -> vec3<f32> {
321    let min_value = min(v.x, min(v.y, v.z));
322    let min_vec = vec3<f32>(min_value);
324    return select(
325        vec3<f32>(0.0),
326        vec3<f32>(1.0),
327        v == min_vec,
328    );
332@vertex
333fn vs_points(
334    @builtin(instance_index) instance_index: u32,
335    @builtin(vertex_index) vertex_index: u32,
336) -> GaussianVertexOutput {
337    var output: GaussianVertexOutput;
339    let entry = get_entry(instance_index);
340    let splat_index = entry.value;
342    var discard_quad = false;
344    discard_quad |= entry.key == 0xFFFFFFFFu; // || splat_index == 0u;
346    let position = vec4<f32>(get_position(splat_index), 1.0);
348    let transformed_position = (gaussian_uniforms.transform * position).xyz;
349    let projected_position = world_to_clip(transformed_position);
351    discard_quad |= !in_frustum(projected_position.xyz);
353#ifdef DRAW_SELECTED
354    discard_quad |= get_visibility(splat_index) < 0.5;
355#endif
357    if (discard_quad) {
358        output.color = vec4<f32>(0.0, 0.0, 0.0, 0.0);
359        output.position = vec4<f32>(0.0, 0.0, 0.0, 0.0);
360        return output;
361    }
363    var quad_vertices = array<vec2<f32>, 4>(
364        vec2<f32>(-1.0, -1.0),
365        vec2<f32>(-1.0,  1.0),
366        vec2<f32>( 1.0, -1.0),
367        vec2<f32>( 1.0,  1.0),
368    );
370    let quad_index = vertex_index % 4u;
371    let quad_offset = quad_vertices[quad_index];
373    let ray_direction = normalize(transformed_position - view.world_position);
375    var rgb = vec3<f32>(0.0);
377#ifdef RASTERIZE_DEPTH
378    // TODO: unbiased depth rendering, see: https://zju3dv.github.io/pgsr/
379    let first_position = vec4<f32>(get_position(get_entry(1u).value), 1.0);
380    let last_position = vec4<f32>(get_position(get_entry(gaussian_uniforms.count - 1u).value), 1.0);
382    let min_position = (gaussian_uniforms.transform * first_position).xyz;
383    let max_position = (gaussian_uniforms.transform * last_position).xyz;
385    let camera_position = view.world_position;
387    let min_distance = length(min_position - camera_position);
388    let max_distance = length(max_position - camera_position);
390    let depth = length(transformed_position - camera_position);
391    rgb = depth_to_rgb(
392        depth,
393        min_distance,
394        max_distance,
395    );
396#else ifdef RASTERIZE_NORMAL
397    let R = get_rotation_matrix(get_rotation(splat_index));
398    let S = get_scale_matrix(get_scale(splat_index));
399    let T = mat3x3<f32>(
400        gaussian_uniforms.transform[0].xyz,
401        gaussian_uniforms.transform[1].xyz,
402        gaussian_uniforms.transform[2].xyz,
403    );
404    let L = T * S * R;
406    let local_normal = vec4<f32>(L.z, 0.0);
407    let world_normal = view.view_from_world * local_normal;
409    let t = normalize(world_normal);
411    rgb = vec3<f32>(
412        0.5 * (t.x + 1.0),
413        0.5 * (t.y + 1.0),
414        0.5 * (t.z + 1.0)
415    );
416#else
417    rgb = get_color(splat_index, ray_direction);
418#endif
420    let opacity = get_opacity(splat_index);
422#ifdef OPACITY_ADAPTIVE_RADIUS
423    let cutoff = sqrt(max(9.0 + 2.0 * log(opacity), 0.000001));
424#else
425    let cutoff = 3.0;
426#endif
428    // TODO: verify color benefit for ray_direction computed at quad verticies instead of gaussian center (same as current complexity)
429    output.color = vec4<f32>(
430        rgb,
431        opacity * gaussian_uniforms.global_opacity,
432    );
434#ifdef HIGHLIGHT_SELECTED
435    if (get_visibility(splat_index) > 0.5) {
436        output.color = vec4<f32>(0.3, 1.0, 0.1, 1.0);
437    }
438#endif
440#ifdef GAUSSIAN_3D
441    let cov2d = compute_cov2d_3dgs(
442        transformed_position,
443        splat_index,
444    );
445    let bb = get_bounding_box(
446        cov2d,
447        quad_offset,
448        cutoff,
449    );
451#ifdef USE_AABB
452    let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y;
453    let det_inv = 1.0 / det;
454    let conic = vec3<f32>(
455        cov2d.z * det_inv,
456        -cov2d.y * det_inv,
457        cov2d.x * det_inv
458    );
459    // TODO: this conic seems only valid in 3dgs
460    output.conic = conic;
461    output.major_minor = bb.zw;
462#endif
464#else ifdef GAUSSIAN_SURFEL
465    let surfel = compute_cov2d_surfel(
466        transformed_position,
467        splat_index,
468        cutoff,
469    );
471    output.local_to_pixel_u = surfel.local_to_pixel.x;
472    output.local_to_pixel_v = surfel.local_to_pixel.y;
473    output.local_to_pixel_w = surfel.local_to_pixel.z;
474    output.mean_2d = surfel.mean_2d;
476    let bb = get_bounding_box_cov2d(
477        surfel.extent,
478        quad_offset,
479        cutoff,
480    );
481    output.radius = bb.zw;
482#endif
484    output.uv = quad_offset;
485    output.position = vec4<f32>(
486        projected_position.xy + bb.xy,
487        projected_position.zw,
488    );
490    return output;
493@fragment
494fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4<f32> {
495#ifdef USE_AABB
496#ifdef GAUSSIAN_SURFEL
497    let radius = input.radius;
498    let mean_2d = input.mean_2d;
499    let aspect = vec2<f32>(
500        1.0,
501        view.viewport.z / view.viewport.w,
502    );
503    let pixel_coord = input.uv * radius * aspect + mean_2d;
504    // let pixel_coord = input.position.xy * view.viewport.zw + view.viewport.xy;
506    let power = surfel_fragment_power(
507        mat3x3<f32>(
508            input.local_to_pixel_u,
509            input.local_to_pixel_v,
510            input.local_to_pixel_w,
511        ),
512        pixel_coord,
513        mean_2d,
514    );
515#else ifdef GAUSSIAN_3D
516    let d = -input.major_minor;
517    let conic = input.conic;
518    let power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y;
519#endif
521    if (power > 0.0) {
522        discard;
523    }
524#endif
526#ifdef USE_OBB
527    let sigma = 1.0 / 3.0;
528    let sigma_squared = 2.0 * sigma * sigma;
529    let distance_squared = dot(input.uv, input.uv);
531    let power = -distance_squared / sigma_squared;
533    if (distance_squared > 3.0 * 3.0) {
534        discard;
535    }
536#endif
538#ifdef VISUALIZE_BOUNDING_BOX
539    let uv = input.uv * 0.5 + 0.5;
540    let edge_width = 0.08;
541    if (
542        (uv.x < edge_width || uv.x > 1.0 - edge_width) ||
543        (uv.y < edge_width || uv.y > 1.0 - edge_width)
544    ) {
545        return vec4<f32>(0.3, 1.0, 0.1, 1.0);
546    }
547#endif
549    let alpha = min(exp(power) * input.color.a, 0.999);
551    // TODO: round alpha to terminate depth test?
553    return vec4<f32>(
554        input.color.rgb * alpha,
555        alpha,
556    );
Contents