1#define_import_path bevy_gaussian_splatting::gaussian_2d
2
3
4#ifdef GAUSSIAN_2D
5#import bevy_gaussian_splatting::bindings::{
6    view,
7    gaussian_uniforms,
8}
9#import bevy_gaussian_splatting::helpers::{
10    get_rotation_matrix,
11    get_scale_matrix,
12    intrinsic_matrix,
13}
14
15#ifdef PACKED
16    #import bevy_gaussian_splatting::packed::{
17        get_position,
18        get_color,
19        get_visibility,
20        get_opacity,
21        get_rotation,
22        get_scale,
23    }
24#else ifdef BUFFER_STORAGE
25    #import bevy_gaussian_splatting::planar::{
26        get_position,
27        get_color,
28        get_visibility,
29        get_opacity,
30        get_rotation,
31        get_scale,
32    }
33#else BUFFER_TEXTURE
34    #import bevy_gaussian_splatting::texture::{
35        get_position,
36        get_color,
37        get_visibility,
38        get_opacity,
39        get_rotation,
40        get_scale,
41    }
42#endif
43
44
45struct Surfel {
46    local_to_pixel: mat3x3<f32>,
47    mean_2d: vec2<f32>,
48    extent: vec2<f32>,
49};
50
51
52fn get_bounding_box_cov2d(
53    extent: vec2<f32>,
54    direction: vec2<f32>,
55    cutoff: f32,
56) -> vec4<f32> {
57    let fitler_size = 0.707106;
58
59    if extent.x < 1.e-4 || extent.y < 1.e-4 {
60        return vec4<f32>(0.0);
61    }
62
63    let radius = sqrt(extent);
64    let max_radius = vec2<f32>(max(
65        max(radius.x, radius.y),
66        cutoff * fitler_size,
67    ));
68
69    // TODO: verify OBB capability
70    let radius_ndc = vec2<f32>(
71        max_radius / view.viewport.zw,
72    );
73
74    return vec4<f32>(
75        radius_ndc * direction,
76        max_radius,
77    );
78}
79
80
81fn compute_cov2d_surfel(
82    gaussian_position: vec3<f32>,
83    index: u32,
84    cutoff: f32,
85) -> Surfel {
86    var output: Surfel;
87
88    let rotation = get_rotation(index);
89    let scale = get_scale(index);
90
91    let T_r = mat3x3<f32>(
92        gaussian_uniforms.transform[0].xyz,
93        gaussian_uniforms.transform[1].xyz,
94        gaussian_uniforms.transform[2].xyz,
95    );
96
97    let S = get_scale_matrix(scale);
98    let R = get_rotation_matrix(rotation);
99
100    let L = T_r * transpose(R) * S;
101
102    let world_from_local = mat3x4<f32>(
103        vec4<f32>(L[0], 0.0),
104        vec4<f32>(L[1], 0.0),
105        vec4<f32>(gaussian_position, 1.0),
106    );
107
108    let ndc_from_world = transpose(view.clip_from_world);
109    let pixels_from_ndc = intrinsic_matrix();
110
111    let T = transpose(world_from_local) * ndc_from_world * pixels_from_ndc;
112
113    let test = vec3<f32>(cutoff * cutoff, cutoff * cutoff, -1.0);
114    let d = dot(test * T[2], T[2]);
115    if abs(d) < 1.0e-4 {
116        output.extent = vec2<f32>(0.0);
117        return output;
118    }
119
120    let f = (1.0 / d) * test;
121    let mean_2d = vec2<f32>(
122        dot(f, T[0] * T[2]),
123        dot(f, T[1] * T[2]),
124    );
125
126    let t = vec2<f32>(
127        dot(f * T[0], T[0]),
128        dot(f * T[1], T[1]),
129    );
130    let extent = mean_2d * mean_2d - t;
131
132    output.local_to_pixel = T;
133    output.mean_2d = mean_2d;
134    output.extent = extent;
135    return output;
136}
137
138fn surfel_fragment_power(
139    local_to_pixel: mat3x3<f32>,
140    pixel_coord: vec2<f32>,
141    mean_2d: vec2<f32>,
142) -> f32 {
143    let deltas = mean_2d - pixel_coord;
144
145    let hu = pixel_coord.x * local_to_pixel[2] - local_to_pixel[0];
146    let hv = pixel_coord.y * local_to_pixel[2] - local_to_pixel[1];
147
148    let p = cross(hu, hv);
149
150    let us = p.x / p.z;
151    let vs = p.y / p.z;
152
153    let sigmas_3d = us * us + vs * vs;
154    let sigmas_2d = 2.0 * (deltas.x * deltas.x + deltas.y * deltas.y);
155
156    let sigmas = 0.5 * min(sigmas_3d, sigmas_2d);
157    let power = -sigmas;
158
159    return power;
160}
161
162#endif  // GAUSSIAN_2D