1#define_import_path bevy_gaussian_splatting::helpers
2
3#import bevy_gaussian_splatting::bindings::{
4    view,
5    gaussian_uniforms,
6}
7
8fn cov2d(
9    position: vec3<f32>,
10    cov3d: array<f32, 6>,
11) -> vec3<f32> {
12    let Vrk = mat3x3(
13        cov3d[0], cov3d[1], cov3d[2],
14        cov3d[1], cov3d[3], cov3d[4],
15        cov3d[2], cov3d[4], cov3d[5],
16    );
17
18    var t = view.view_from_world * vec4<f32>(position, 1.0);
19
20    let focal = vec2<f32>(
21        view.clip_from_view[0].x * view.viewport.z,
22        view.clip_from_view[1].y * view.viewport.w,
23    );
24
25    let s = 1.0 / (t.z * t.z);
26    let J = mat3x3(
27        focal.x / t.z, 0.0, -(focal.x * t.x) * s,
28        0.0, -focal.y / t.z, (focal.y * t.y) * s,
29        0.0, 0.0, 0.0,
30    );
31
32    let W = transpose(
33        mat3x3<f32>(
34            view.view_from_world[0].xyz,
35            view.view_from_world[1].xyz,
36            view.view_from_world[2].xyz,
37        )
38    );
39
40    let T = W * J;
41
42    var cov = transpose(T) * transpose(Vrk) * T;
43    cov[0][0] += 0.3f;
44    cov[1][1] += 0.3f;
45
46    return vec3<f32>(cov[0][0], cov[0][1], cov[1][1]);
47}
48
49fn get_bounding_box_clip(
50    cov2d: vec3<f32>,
51    direction: vec2<f32>,
52    cutoff: f32,
53) -> vec4<f32> {
54    // return vec4<f32>(offset, uv);
55
56    let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y;
57    let trace = cov2d.x + cov2d.z;
58    let mid = 0.5 * trace;
59    let discriminant = max(0.0, mid * mid - det);
60
61    let term = sqrt(discriminant);
62
63    let lambda1 = mid + term;
64    let lambda2 = max(mid - term, 0.0);
65
66    let x_axis_length = sqrt(lambda1);
67    let y_axis_length = sqrt(lambda2);
68
69#ifdef USE_AABB
70    let radius_px = cutoff * max(x_axis_length, y_axis_length);
71    let radius_ndc = vec2<f32>(
72        radius_px / view.viewport.zw,
73    );
74
75    return vec4<f32>(
76        radius_ndc * direction,
77        radius_px * direction,
78    );
79#endif
80
81#ifdef USE_OBB
82
83    let a = (cov2d.x - cov2d.z) * (cov2d.x - cov2d.z);
84    let b = sqrt(a + 4.0 * cov2d.y * cov2d.y);
85    let major_radius = sqrt((cov2d.x + cov2d.z + b) * 0.5);
86    let minor_radius = sqrt((cov2d.x + cov2d.z - b) * 0.5);
87
88    let bounds = cutoff * vec2<f32>(
89        major_radius,
90        minor_radius,
91    );
92
93    let eigvec1 = normalize(vec2<f32>(
94        -cov2d.y,
95        lambda1 - cov2d.x,
96    ));
97    let eigvec2 = vec2<f32>(
98        eigvec1.y,
99        -eigvec1.x
100    );
101
102    let rotation_matrix = transpose(
103        mat2x2(
104            eigvec1,
105            eigvec2,
106        )
107    );
108
109    let scaled_vertex = direction * bounds;
110    let rotated_vertex = scaled_vertex * rotation_matrix;
111
112    let scaling_factor = 1.0 / view.viewport.zw;
113    let ndc_vertex = rotated_vertex * scaling_factor;
114
115    return vec4<f32>(
116        ndc_vertex,
117        rotated_vertex,
118    );
119#endif
120}
121
122fn intrinsic_matrix() -> mat3x4<f32> {
123    let focal = vec2<f32>(
124        view.clip_from_view[0].x * view.viewport.z / 2.0,
125        view.clip_from_view[1].y * view.viewport.w / 2.0,
126    );
127
128    let Ks = mat3x4<f32>(
129        vec4<f32>(focal.x, 0.0, 0.0, (view.viewport.z - 1.0) / 2.0),
130        vec4<f32>(0.0, focal.y, 0.0, (view.viewport.w - 1.0) / 2.0),
131        vec4<f32>(0.0, 0.0, 0.0, 1.0)
132    );
133
134    return Ks;
135}
136
137fn get_rotation_matrix(
138    rotation: vec4<f32>,
139) -> mat3x3<f32> {
140    let r = rotation.x;
141    let x = rotation.y;
142    let y = rotation.z;
143    let z = rotation.w;
144
145    return mat3x3<f32>(
146        1.0 - 2.0 * (y * y + z * z),
147        2.0 * (x * y - r * z),
148        2.0 * (x * z + r * y),
149
150        2.0 * (x * y + r * z),
151        1.0 - 2.0 * (x * x + z * z),
152        2.0 * (y * z - r * x),
153
154        2.0 * (x * z - r * y),
155        2.0 * (y * z + r * x),
156        1.0 - 2.0 * (x * x + y * y),
157    );
158}
159
160fn get_scale_matrix(
161    scale: vec3<f32>,
162) -> mat3x3<f32> {
163    return mat3x3<f32>(
164        scale.x * gaussian_uniforms.global_scale, 0.0, 0.0,
165        0.0, scale.y * gaussian_uniforms.global_scale, 0.0,
166        0.0, 0.0, scale.z * gaussian_uniforms.global_scale,
167    );
168}