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