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    Entry,
5}
6#import bevy_gaussian_splatting::classification::class_to_rgb
7#import bevy_gaussian_splatting::depth::depth_to_rgb
8#import bevy_gaussian_splatting::optical_flow::{
9    calculate_motion_vector,
10    optical_flow_to_rgb,
12#import bevy_gaussian_splatting::helpers::{
13    get_rotation_matrix,
14    get_scale_matrix,
16#import bevy_gaussian_splatting::transform::{
17    world_to_clip,
18    in_frustum,
21#ifdef GAUSSIAN_2D
22    #import bevy_gaussian_splatting::gaussian_2d::{
23        compute_cov2d_surfel,
24        get_bounding_box_cov2d,
25        surfel_fragment_power,
26    }
27#else ifdef GAUSSIAN_3D
28    #import bevy_gaussian_splatting::gaussian_3d::{
29        compute_cov2d_3dgs,
30    }
31    #import bevy_gaussian_splatting::helpers::{
32        get_bounding_box_clip,
33    }
34#else ifdef GAUSSIAN_4D
35    #import bevy_gaussian_splatting::gaussian_4d::{
36        conditional_cov3d,
37    }
38    #import bevy_gaussian_splatting::helpers::{
39        cov2d,
40        get_bounding_box_clip,
41    }
42#endif
44#ifdef PACKED
45    #ifdef PRECOMPUTE_COVARIANCE_3D
46        #import bevy_gaussian_splatting::packed::{
47            get_position,
48            get_color,
49            get_visibility,
50            get_opacity,
51            get_cov3d,
52        }
53    #else
54        #import bevy_gaussian_splatting::packed::{
55            get_position,
56            get_color,
57            get_visibility,
58            get_opacity,
59            get_rotation,
60            get_scale,
61        }
62    #endif
63#else ifdef BUFFER_STORAGE
64    #ifdef PRECOMPUTE_COVARIANCE_3D
65        #import bevy_gaussian_splatting::planar::{
66            get_position,
67            get_color,
68            get_visibility,
69            get_opacity,
70            get_cov3d,
71        }
72    #else
73        #import bevy_gaussian_splatting::planar::{
74            get_position,
75            get_color,
76            get_visibility,
77            get_opacity,
78            get_rotation,
79            get_scale,
80        }
81    #endif
82#else ifdef BUFFER_TEXTURE
83    #ifdef PRECOMPUTE_COVARIANCE_3D
84        #import bevy_gaussian_splatting::texture::{
85            get_position,
86            get_color,
87            get_visibility,
88            get_opacity,
89            get_cov3d,
90            location,
91        }
92    #else
93        #import bevy_gaussian_splatting::texture::{
94            get_position,
95            get_color,
96            get_visibility,
97            get_opacity,
98            get_rotation,
99            get_scale,
100            location,
101        }
102    #endif
103#endif
106#ifdef BUFFER_STORAGE
107    @group(3) @binding(0) var<storage, read> sorted_entries: array<Entry>;
108    fn get_entry(index: u32) -> Entry {
109        return sorted_entries[index];
110    }
111#else ifdef BUFFER_TEXTURE
112    @group(3) @binding(0) var sorted_entries: texture_2d<u32>;
113    fn get_entry(index: u32) -> Entry {
114        let sample = textureLoad(
115            sorted_entries,
116            location(index),
117            0,
118        );
120        return Entry(
121            sample.r,
122            sample.g,
123        );
124    }
125#endif
127#ifdef WEBGL2
128    struct GaussianVertexOutput {
129        @builtin(position) position: vec4<f32>,
130        @location(0) color: vec4<f32>,
131        @location(1) uv: vec2<f32>,
132    #ifdef GAUSSIAN_2D
133        @location(2) local_to_pixel_u: vec3<f32>,
134        @location(3) local_to_pixel_v: vec3<f32>,
135        @location(4) local_to_pixel_w: vec3<f32>,
136        @location(5) mean_2d: vec2<f32>,
137        @location(6) radius: vec2<f32>,
138    #else #ifdef GAUSSIAN_3D
139        @location(2) conic: vec3<f32>,
140        @location(3) major_minor: vec2<f32>,
141    #else #ifdef GAUSSIAN_4D
142        @location(2) conic: vec3<f32>,
143        @location(3) major_minor: vec2<f32>,
144    #endif
145    };
146#else
147    struct GaussianVertexOutput {
148        @builtin(position) position: vec4<f32>,
149        @location(0) @interpolate(flat) color: vec4<f32>,
150        @location(1) @interpolate(linear) uv: vec2<f32>,
151    #ifdef GAUSSIAN_2D
152        @location(2) @interpolate(flat) local_to_pixel_u: vec3<f32>,
153        @location(3) @interpolate(flat) local_to_pixel_v: vec3<f32>,
154        @location(4) @interpolate(flat) local_to_pixel_w: vec3<f32>,
155        @location(5) @interpolate(flat) mean_2d: vec2<f32>,
156        @location(6) @interpolate(flat) radius: vec2<f32>,
157    #else ifdef GAUSSIAN_3D
158        @location(2) @interpolate(flat) conic: vec3<f32>,
159        @location(3) @interpolate(linear) major_minor: vec2<f32>,
160    #else ifdef GAUSSIAN_4D
161        @location(2) @interpolate(flat) conic: vec3<f32>,
162        @location(3) @interpolate(linear) major_minor: vec2<f32>,
163    #endif
164    };
165#endif
168@vertex
169fn vs_points(
170    @builtin(instance_index) instance_index: u32,
171    @builtin(vertex_index) vertex_index: u32,
172) -> GaussianVertexOutput {
173    var output: GaussianVertexOutput;
175    let entry = get_entry(instance_index);
176    let splat_index = entry.value;
178    var discard_quad = false;
180    discard_quad |= entry.key == 0xFFFFFFFFu; // || splat_index == 0u;
182    let position = vec4<f32>(get_position(splat_index), 1.0);
184    var transformed_position = (gaussian_uniforms.transform * position).xyz;
185    var previous_transformed_position = transformed_position;
187#ifdef DRAW_SELECTED
188    discard_quad |= get_visibility(splat_index) < 0.5;
189#endif
191#ifdef GAUSSIAN_4D
192#else
193    let projected_position = world_to_clip(transformed_position);
194    discard_quad |= !in_frustum(projected_position.xyz);
195#endif
197    if (discard_quad) {
198        output.color = vec4<f32>(0.0, 0.0, 0.0, 0.0);
199        output.position = vec4<f32>(0.0, 0.0, 0.0, 0.0);
200        return output;
201    }
203    var quad_vertices = array<vec2<f32>, 4>(
204        vec2<f32>(-1.0, -1.0),
205        vec2<f32>(-1.0,  1.0),
206        vec2<f32>( 1.0, -1.0),
207        vec2<f32>( 1.0,  1.0),
208    );
210    let quad_index = vertex_index % 4u;
211    let quad_offset = quad_vertices[quad_index];
214    var opacity = get_opacity(splat_index);
216#ifdef OPACITY_ADAPTIVE_RADIUS
217    let cutoff = sqrt(max(9.0 + 2.0 * log(opacity), 0.000001));
218#else
219    let cutoff = 3.0;
220#endif
223#ifdef GAUSSIAN_2D
224    let surfel = compute_cov2d_surfel(
225        transformed_position,
226        splat_index,
227        cutoff,
228    );
230    output.local_to_pixel_u = surfel.local_to_pixel[0];
231    output.local_to_pixel_v = surfel.local_to_pixel[1];
232    output.local_to_pixel_w = surfel.local_to_pixel[2];
233    output.mean_2d = surfel.mean_2d;
235    let bb = get_bounding_box_cov2d(
236        surfel.extent,
237        quad_offset,
238        cutoff,
239    );
240    output.radius = bb.zw;
241#else
242    #ifdef GAUSSIAN_3D
243        let gaussian_cov2d = compute_cov2d_3dgs(
244            transformed_position,
245            splat_index,
246        );
247    #else ifdef GAUSSIAN_4D
248        let gaussian_4d = conditional_cov3d(
249            transformed_position,
250            splat_index,
251            gaussian_uniforms.time,
252        );
254        if !gaussian_4d.mask {
255            output.color = vec4<f32>(0.0, 0.0, 0.0, 0.0);
256            output.position = vec4<f32>(0.0, 0.0, 0.0, 0.0);
257            return output;
258        }
260        let position_t = vec4<f32>(position.xyz + gaussian_4d.delta_mean, 1.0);
261        transformed_position = (gaussian_uniforms.transform * position_t).xyz;
262        // TODO: set previous_transformed_position based on temporal position delta
263        let projected_position = world_to_clip(transformed_position);
265        if !in_frustum(projected_position.xyz) {
266            output.color = vec4<f32>(0.0, 0.0, 0.0, 0.0);
267            output.position = vec4<f32>(0.0, 0.0, 0.0, 0.0);
268            return output;
269        }
271        opacity = opacity * gaussian_4d.opacity_modifier;
273        let gaussian_cov2d = cov2d(
274            transformed_position,
275            gaussian_4d.cov3d,
276        );
277    #endif
279    let bb = get_bounding_box_clip(
280        gaussian_cov2d,
281        quad_offset,
282        cutoff,
283    );
285    #ifdef USE_AABB
286        let det = gaussian_cov2d.x * gaussian_cov2d.z - gaussian_cov2d.y * gaussian_cov2d.y;
287        let det_inv = 1.0 / det;
288        let conic = vec3<f32>(
289            gaussian_cov2d.z * det_inv,
290            -gaussian_cov2d.y * det_inv,
291            gaussian_cov2d.x * det_inv
292        );
293        output.conic = conic;
294        output.major_minor = bb.zw;
295    #endif
296#endif
299    var rgb = vec3<f32>(0.0);
301// TODO: RASTERIZE_ACCELERATION
302#ifdef RASTERIZE_CLASSIFICATION
303    let ray_direction = normalize(transformed_position - view.world_position);
305    #ifdef GAUSSIAN_3D_STRUCTURE
306        rgb = get_color(splat_index, ray_direction);
307    #else ifdef GAUSSIAN_4D
308        rgb = get_color(splat_index, gaussian_4d.dir_t, ray_direction);
309    #endif
311    rgb = class_to_rgb(
312        get_visibility(splat_index),
313        rgb,
314    );
315#else ifdef RASTERIZE_DEPTH
316    // TODO: unbiased depth rendering, see: https://zju3dv.github.io/pgsr/
317    let first_position = vec4<f32>(get_position(get_entry(1u).value), 1.0);
318    let last_position = vec4<f32>(get_position(get_entry(gaussian_uniforms.count - 1u).value), 1.0);
320    let min_position = (gaussian_uniforms.transform * first_position).xyz;
321    let max_position = (gaussian_uniforms.transform * last_position).xyz;
323    let camera_position = view.world_position;
325    let min_distance = length(min_position - camera_position);
326    let max_distance = length(max_position - camera_position);
328    let depth = length(transformed_position - camera_position);
329    rgb = depth_to_rgb(
330        depth,
331        min_distance,
332        max_distance,
333    );
334#else ifdef RASTERIZE_NORMAL
335    // TODO: support rotation decomposition for 4d gaussians
336    let R = get_rotation_matrix(get_rotation(splat_index));
337    let S = get_scale_matrix(get_scale(splat_index));
338    let T = mat3x3<f32>(
339        gaussian_uniforms.transform[0].xyz,
340        gaussian_uniforms.transform[1].xyz,
341        gaussian_uniforms.transform[2].xyz,
342    );
343    let L = T * S * R;
345    let local_normal = vec4<f32>(L[2], 0.0);
346    let world_normal = view.view_from_world * local_normal;
348    let t = normalize(world_normal);
350    rgb = vec3<f32>(
351        0.5 * (t.x + 1.0),
352        0.5 * (t.y + 1.0),
353        0.5 * (t.z + 1.0)
354    );
355#else ifdef RASTERIZE_OPTICAL_FLOW
356    let motion_vector = calculate_motion_vector(
357        transformed_position,
358        previous_transformed_position,
359    );
361    rgb = optical_flow_to_rgb(motion_vector);
362#else ifdef RASTERIZE_POSITION
363    rgb = (transformed_position - gaussian_uniforms.min.xyz) / (gaussian_uniforms.max.xyz - gaussian_uniforms.min.xyz);
364#else ifdef RASTERIZE_VELOCITY
365    let time_delta = 1e-3;
366    let future_gaussian_4d = conditional_cov3d(
367        transformed_position,
368        splat_index,
369        gaussian_uniforms.time + time_delta,
370    );
371    let position_delta = future_gaussian_4d.delta_mean - gaussian_4d.delta_mean;
372    let velocity = position_delta / time_delta;
373    let velocity_magnitude = length(velocity);
374    let velocity_normalized = normalize(velocity);
376    // TODO: magnitude normalization
377    let min_magnitude = 1.0;
378    let max_magnitude = 2.0;
380    let scaled_mag = clamp(
381        (velocity_magnitude - min_magnitude) / (max_magnitude - min_magnitude),
382        0.0,
383        1.0
384    );
386    if scaled_mag < 1e-2 {
387        opacity = 0.0;
388    }
390    let base_color = 0.5 * (velocity_normalized + vec3<f32>(1.0, 1.0, 1.0));
391    rgb = base_color * scaled_mag;
392#else ifdef RASTERIZE_COLOR
393    // TODO: verify color benefit for ray_direction computed at quad verticies instead of gaussian center (same as current complexity)
394    // TODO: why doesn't Transform rotation change SH color?
395    let ray_direction = normalize(transformed_position - view.world_position);
397    #ifdef GAUSSIAN_3D_STRUCTURE
398        rgb = get_color(splat_index, ray_direction);
399    #else ifdef GAUSSIAN_4D
400        rgb = get_color(splat_index, gaussian_4d.dir_t, ray_direction);
401    #endif
402#endif
405    output.color = vec4<f32>(
406        rgb,
407        opacity * gaussian_uniforms.global_opacity,
408    );
410#ifdef HIGHLIGHT_SELECTED
411    if (get_visibility(splat_index) > 0.5) {
412        output.color = vec4<f32>(0.3, 1.0, 0.1, 1.0);
413    }
414#endif
417    output.uv = quad_offset;
418    output.position = vec4<f32>(
419        projected_position.xy + bb.xy,
420        projected_position.zw,
421    );
423    return output;
426@fragment
427fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4<f32> {
428#ifdef USE_AABB
429#ifdef GAUSSIAN_2D
430    let radius = input.radius;
431    let mean_2d = input.mean_2d;
432    let aspect = vec2<f32>(
433        1.0,
434        view.viewport.z / view.viewport.w,
435    );
436    let pixel_coord = input.uv * radius * aspect + mean_2d;
438    let power = surfel_fragment_power(
439        mat3x3<f32>(
440            input.local_to_pixel_u,
441            input.local_to_pixel_v,
442            input.local_to_pixel_w,
443        ),
444        pixel_coord,
445        mean_2d,
446    );
447#else ifdef GAUSSIAN_3D
448    let d = -input.major_minor;
449    let conic = input.conic;
450    let power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y;
451#else ifdef GAUSSIAN_4D
452    let d = -input.major_minor;
453    let conic = input.conic;
454    let power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y;
455#endif
457    if (power > 0.0) {
458        discard;
459    }
460#endif
462#ifdef USE_OBB
463    let sigma = 1.0 / 3.0;
464    let sigma_squared = 2.0 * sigma * sigma;
465    let distance_squared = dot(input.uv, input.uv);
467    let power = -distance_squared / sigma_squared;
469    if (distance_squared > 3.0 * 3.0) {
470        discard;
471    }
472#endif
474#ifdef VISUALIZE_BOUNDING_BOX
475    let uv = input.uv * 0.5 + 0.5;
476    let edge_width = 0.08;
477    if (
478        (uv.x < edge_width || uv.x > 1.0 - edge_width) ||
479        (uv.y < edge_width || uv.y > 1.0 - edge_width)
480    ) {
481        return vec4<f32>(0.3, 1.0, 0.1, 1.0);
482    }
483#endif
485    let alpha = min(exp(power) * input.color.a, 0.999);
487    // TODO: round alpha to terminate depth test?
489    return vec4<f32>(
490        input.color.rgb * alpha,
491        alpha,
492    );
Contents