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