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