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
105#ifdef BUFFER_STORAGE
106    @group(3) @binding(0) var<storage, read> sorted_entries: array<Entry>;
107    fn get_entry(index: u32) -> Entry {
108        return sorted_entries[index];
109    }
110#else ifdef BUFFER_TEXTURE
111    @group(3) @binding(0) var sorted_entries: texture_2d<u32>;
112    fn get_entry(index: u32) -> Entry {
113        let sample = textureLoad(
114            sorted_entries,
115            location(index),
116            0,
117        );
119        return Entry(
120            sample.r,
121            sample.g,
122        );
123    }
124#endif
126#ifdef WEBGL2
127    struct GaussianVertexOutput {
128        @builtin(position) position: vec4<f32>,
129        @location(0) color: vec4<f32>,
130        @location(1) uv: vec2<f32>,
131    #ifdef GAUSSIAN_2D
132        @location(2) local_to_pixel_u: vec3<f32>,
133        @location(3) local_to_pixel_v: vec3<f32>,
134        @location(4) local_to_pixel_w: vec3<f32>,
135        @location(5) mean_2d: vec2<f32>,
136        @location(6) radius: vec2<f32>,
137    #else #ifdef GAUSSIAN_3D
138        @location(2) conic: vec3<f32>,
139        @location(3) major_minor: vec2<f32>,
140    #else #ifdef GAUSSIAN_4D
141        @location(2) conic: vec3<f32>,
142        @location(3) major_minor: vec2<f32>,
143    #endif
144    };
145#else
146    struct GaussianVertexOutput {
147        @builtin(position) position: vec4<f32>,
148        @location(0) @interpolate(flat) color: vec4<f32>,
149        @location(1) @interpolate(linear) uv: vec2<f32>,
150    #ifdef GAUSSIAN_2D
151        @location(2) @interpolate(flat) local_to_pixel_u: vec3<f32>,
152        @location(3) @interpolate(flat) local_to_pixel_v: vec3<f32>,
153        @location(4) @interpolate(flat) local_to_pixel_w: vec3<f32>,
154        @location(5) @interpolate(flat) mean_2d: vec2<f32>,
155        @location(6) @interpolate(flat) radius: vec2<f32>,
156    #else ifdef GAUSSIAN_3D
157        @location(2) @interpolate(flat) conic: vec3<f32>,
158        @location(3) @interpolate(linear) major_minor: vec2<f32>,
159    #else ifdef GAUSSIAN_4D
160        @location(2) @interpolate(flat) conic: vec3<f32>,
161        @location(3) @interpolate(linear) major_minor: vec2<f32>,
162    #endif
163    };
164#endif
166fn world_to_local_direction(ray_direction_world: vec3<f32>, transform: mat4x4<f32>) -> vec3<f32> {
167    let basis = mat3x3<f32>(
168        transform[0].xyz,
169        transform[1].xyz,
170        transform[2].xyz,
171    );
172    let basis_x = normalize(basis[0]);
173    let basis_y = normalize(basis[1]);
174    let basis_z = normalize(basis[2]);
176    let local = vec3<f32>(
177        dot(basis_x, ray_direction_world),
178        dot(basis_y, ray_direction_world),
179        dot(basis_z, ray_direction_world),
180    );
182    return normalize(local);
184@vertex
185fn vs_points(
186    @builtin(instance_index) instance_index: u32,
187    @builtin(vertex_index) vertex_index: u32,
188) -> GaussianVertexOutput {
189    var output: GaussianVertexOutput;
191    let entry = get_entry(instance_index);
192    let splat_index = entry.value;
194    var discard_quad = false;
196    discard_quad |= entry.key == 0xFFFFFFFFu; // || splat_index == 0u;
198    let position = vec4<f32>(get_position(splat_index), 1.0);
200    var transformed_position = (gaussian_uniforms.transform * position).xyz;
201    var previous_transformed_position = transformed_position;
203#ifdef DRAW_SELECTED
204    discard_quad |= get_visibility(splat_index) < 0.5;
205#endif
207#ifdef GAUSSIAN_4D
208#else
209    let projected_position = world_to_clip(transformed_position);
210    discard_quad |= !in_frustum(projected_position.xyz);
211#endif
213    if (discard_quad) {
214        output.color = vec4<f32>(0.0, 0.0, 0.0, 0.0);
215        output.position = vec4<f32>(0.0, 0.0, 0.0, 0.0);
216        return output;
217    }
219    var quad_vertices = array<vec2<f32>, 4>(
220        vec2<f32>(-1.0, -1.0),
221        vec2<f32>(-1.0,  1.0),
222        vec2<f32>( 1.0, -1.0),
223        vec2<f32>( 1.0,  1.0),
224    );
226    let quad_index = vertex_index % 4u;
227    let quad_offset = quad_vertices[quad_index];
229    var opacity = get_opacity(splat_index);
231#ifdef OPACITY_ADAPTIVE_RADIUS
232    let cutoff = sqrt(max(9.0 + 2.0 * log(opacity), 0.000001));
233#else
234    let cutoff = 3.0;
235#endif
237#ifdef GAUSSIAN_2D
238    let surfel = compute_cov2d_surfel(
239        transformed_position,
240        splat_index,
241        cutoff,
242    );
244    output.local_to_pixel_u = surfel.local_to_pixel[0];
245    output.local_to_pixel_v = surfel.local_to_pixel[1];
246    output.local_to_pixel_w = surfel.local_to_pixel[2];
247    output.mean_2d = surfel.mean_2d;
249    let bb = get_bounding_box_cov2d(
250        surfel.extent,
251        quad_offset,
252        cutoff,
253    );
254    output.radius = bb.zw;
255#else
256    #ifdef GAUSSIAN_3D
257        let gaussian_cov2d = compute_cov2d_3dgs(
258            transformed_position,
259            splat_index,
260        );
261    #else ifdef GAUSSIAN_4D
262        let gaussian_4d = conditional_cov3d(
263            transformed_position,
264            splat_index,
265            gaussian_uniforms.time,
266        );
268        if !gaussian_4d.mask {
269            output.color = vec4<f32>(0.0, 0.0, 0.0, 0.0);
270            output.position = vec4<f32>(0.0, 0.0, 0.0, 0.0);
271            return output;
272        }
274        let position_t = vec4<f32>(position.xyz + gaussian_4d.delta_mean, 1.0);
275        transformed_position = (gaussian_uniforms.transform * position_t).xyz;
276        // TODO: set previous_transformed_position based on temporal position delta
277        let projected_position = world_to_clip(transformed_position);
279        if !in_frustum(projected_position.xyz) {
280            output.color = vec4<f32>(0.0, 0.0, 0.0, 0.0);
281            output.position = vec4<f32>(0.0, 0.0, 0.0, 0.0);
282            return output;
283        }
285        opacity = opacity * gaussian_4d.opacity_modifier;
287        let gaussian_cov2d = cov2d(
288            transformed_position,
289            gaussian_4d.cov3d,
290        );
291    #endif
293    let bb = get_bounding_box_clip(
294        gaussian_cov2d,
295        quad_offset,
296        cutoff,
297    );
299    #ifdef USE_AABB
300        let det = gaussian_cov2d.x * gaussian_cov2d.z - gaussian_cov2d.y * gaussian_cov2d.y;
301        let det_inv = 1.0 / det;
302        let conic = vec3<f32>(
303            gaussian_cov2d.z * det_inv,
304            -gaussian_cov2d.y * det_inv,
305            gaussian_cov2d.x * det_inv
306        );
307        output.conic = conic;
308        output.major_minor = bb.zw;
309    #endif
310#endif
312    var rgb = vec3<f32>(0.0);
314// TODO: RASTERIZE_ACCELERATION
315#ifdef RASTERIZE_CLASSIFICATION
316    let ray_direction_world = normalize(transformed_position - view.world_position);
317    let ray_direction_local = world_to_local_direction(ray_direction_world, gaussian_uniforms.transform);
319    #ifdef GAUSSIAN_3D_STRUCTURE
320        rgb = get_color(splat_index, ray_direction_local);
321    #else ifdef GAUSSIAN_4D
322        rgb = get_color(splat_index, gaussian_4d.dir_t, ray_direction_local);
323    #endif
325    rgb = class_to_rgb(
326        get_visibility(splat_index),
327        rgb,
328    );
329#else ifdef RASTERIZE_DEPTH
330    // TODO: unbiased depth rendering, see: https://zju3dv.github.io/pgsr/
331    let first_position = vec4<f32>(get_position(get_entry(1u).value), 1.0);
332    let last_position = vec4<f32>(get_position(get_entry(gaussian_uniforms.count - 1u).value), 1.0);
334    let min_position = (gaussian_uniforms.transform * first_position).xyz;
335    let max_position = (gaussian_uniforms.transform * last_position).xyz;
337    let camera_position = view.world_position;
339    let min_distance = length(min_position - camera_position);
340    let max_distance = length(max_position - camera_position);
342    let depth = length(transformed_position - camera_position);
343    rgb = depth_to_rgb(
344        depth,
345        min_distance,
346        max_distance,
347    );
348#else ifdef RASTERIZE_NORMAL
349    // TODO: support rotation decomposition for 4d gaussians
350    let R = get_rotation_matrix(get_rotation(splat_index));
351    let S = get_scale_matrix(get_scale(splat_index));
352    let T = mat3x3<f32>(
353        gaussian_uniforms.transform[0].xyz,
354        gaussian_uniforms.transform[1].xyz,
355        gaussian_uniforms.transform[2].xyz,
356    );
357    let L = T * S * R;
359    let local_normal = vec4<f32>(L[2], 0.0);
360    let world_normal = view.view_from_world * local_normal;
362    let t = normalize(world_normal);
364    rgb = vec3<f32>(
365        0.5 * (t.x + 1.0),
366        0.5 * (t.y + 1.0),
367        0.5 * (t.z + 1.0)
368    );
369#else ifdef RASTERIZE_OPTICAL_FLOW
370    let motion_vector = calculate_motion_vector(
371        transformed_position,
372        previous_transformed_position,
373    );
375    rgb = optical_flow_to_rgb(motion_vector);
376#else ifdef RASTERIZE_POSITION
377    rgb = (transformed_position - gaussian_uniforms.min.xyz) / (gaussian_uniforms.max.xyz - gaussian_uniforms.min.xyz);
378#else ifdef RASTERIZE_VELOCITY
379    let time_delta = 1e-3;
380    let future_gaussian_4d = conditional_cov3d(
381        transformed_position,
382        splat_index,
383        gaussian_uniforms.time + time_delta,
384    );
385    let position_delta = future_gaussian_4d.delta_mean - gaussian_4d.delta_mean;
386    let velocity = position_delta / time_delta;
387    let velocity_magnitude = length(velocity);
388    let velocity_normalized = normalize(velocity);
390    // TODO: magnitude normalization
391    let min_magnitude = 1.0;
392    let max_magnitude = 2.0;
394    let scaled_mag = clamp(
395        (velocity_magnitude - min_magnitude) / (max_magnitude - min_magnitude),
396        0.0,
397        1.0
398    );
400    if scaled_mag < 1e-2 {
401        opacity = 0.0;
402    }
404    let base_color = 0.5 * (velocity_normalized + vec3<f32>(1.0, 1.0, 1.0));
405    rgb = base_color * scaled_mag;
406#else ifdef RASTERIZE_COLOR
407    // TODO: verify color benefit for ray_direction computed at quad verticies instead of gaussian center (same as current complexity)
408    let ray_direction_world = normalize(transformed_position - view.world_position);
409    let ray_direction_local = world_to_local_direction(ray_direction_world, gaussian_uniforms.transform);
411    #ifdef GAUSSIAN_3D_STRUCTURE
412        rgb = get_color(splat_index, ray_direction_local);
413    #else ifdef GAUSSIAN_4D
414        rgb = get_color(splat_index, gaussian_4d.dir_t, ray_direction_local);
415    #endif
416#endif
418    output.color = vec4<f32>(
419        rgb,
420        opacity * gaussian_uniforms.global_opacity,
421    );
423#ifdef HIGHLIGHT_SELECTED
424    if (get_visibility(splat_index) > 0.5) {
425        output.color = vec4<f32>(0.3, 1.0, 0.1, 1.0);
426    }
427#endif
429    output.uv = quad_offset;
430    output.position = vec4<f32>(
431        projected_position.xy + bb.xy,
432        projected_position.zw,
433    );
435    return output;
438@fragment
439fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4<f32> {
440#ifdef USE_AABB
441#ifdef GAUSSIAN_2D
442    let radius = input.radius;
443    let mean_2d = input.mean_2d;
444    let aspect = vec2<f32>(
445        1.0,
446        view.viewport.z / view.viewport.w,
447    );
448    let pixel_coord = input.uv * radius * aspect + mean_2d;
450    let power = surfel_fragment_power(
451        mat3x3<f32>(
452            input.local_to_pixel_u,
453            input.local_to_pixel_v,
454            input.local_to_pixel_w,
455        ),
456        pixel_coord,
457        mean_2d,
458    );
459#else ifdef GAUSSIAN_3D
460    let d = -input.major_minor;
461    let conic = input.conic;
462    let power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y;
463#else ifdef GAUSSIAN_4D
464    let d = -input.major_minor;
465    let conic = input.conic;
466    let power = -0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y;
467#endif
469    if (power > 0.0) {
470        discard;
471    }
472#endif
474#ifdef USE_OBB
475    let sigma = 1.0 / 3.0;
476    let sigma_squared = 2.0 * sigma * sigma;
477    let distance_squared = dot(input.uv, input.uv);
479    let power = -distance_squared / sigma_squared;
481    if (distance_squared > 3.0 * 3.0) {
482        discard;
483    }
484#endif
486#ifdef VISUALIZE_BOUNDING_BOX
487    let uv = input.uv * 0.5 + 0.5;
488    let edge_width = 0.08;
489    if (
490        (uv.x < edge_width || uv.x > 1.0 - edge_width) ||
491        (uv.y < edge_width || uv.y > 1.0 - edge_width)
492    ) {
493        return vec4<f32>(0.3, 1.0, 0.1, 1.0);
494    }
495#endif
497    let alpha = min(exp(power) * input.color.a, 0.999);
499    // TODO: round alpha to terminate depth test?
501    return vec4<f32>(
502        input.color.rgb * alpha,
503        alpha,
504    );
Contents