1#define_import_path bevy_gaussian_splatting::morph::interpolate
2
3#import bevy_gaussian_splatting::bindings::gaussian_uniforms
4
5#ifdef PACKED_F32
6    #import bevy_gaussian_splatting::packed::{
7        get_opacity,
8        get_position,
9        get_rotation,
10        get_scale,
11        get_spherical_harmonics,
12        get_visibility,
13        get_rhs_opacity,
14        get_rhs_position,
15        get_rhs_rotation,
16        get_rhs_scale,
17        get_rhs_spherical_harmonics,
18        get_rhs_visibility,
19        set_output_position_visibility,
20        set_output_spherical_harmonics,
21        set_output_transform,
22    };
23#else
24    #import bevy_gaussian_splatting::planar::{
25        get_opacity,
26        get_position,
27        get_rotation,
28        get_scale,
29        get_spherical_harmonics,
30        get_visibility,
31        get_rhs_opacity,
32        get_rhs_position,
33        get_rhs_rotation,
34        get_rhs_scale,
35        get_rhs_spherical_harmonics,
36        get_rhs_visibility,
37        set_output_position_visibility,
38        set_output_spherical_harmonics,
39        set_output_transform,
40    };
41
42    #ifdef PRECOMPUTE_COVARIANCE_3D
43        #import bevy_gaussian_splatting::planar::{
44            get_cov3d,
45            get_rhs_cov3d,
46            set_output_covariance,
47        };
48    #endif
49#endif
50
51fn interpolation_factor() -> f32 {
52    let duration = gaussian_uniforms.time_stop - gaussian_uniforms.time_start;
53    if abs(duration) < 1e-6 {
54        return select(0.0, 1.0, gaussian_uniforms.time >= gaussian_uniforms.time_stop);
55    }
56    return clamp((gaussian_uniforms.time - gaussian_uniforms.time_start) / duration, 0.0, 1.0);
57}
58
59fn normalize_quaternion(q: vec4<f32>) -> vec4<f32> {
60    let length_squared = dot(q, q);
61    if length_squared <= 0.0 {
62        return vec4<f32>(0.0, 0.0, 0.0, 1.0);
63    }
64    return q / sqrt(length_squared);
65}
66
67const WORKGROUP_SIZE: u32 = 256u;
68
69@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
70fn interpolate_gaussians(
71    @builtin(global_invocation_id) global_id: vec3<u32>,
72) {
73    let index = global_id.x;
74    if index >= gaussian_uniforms.count {
75        return;
76    }
77
78    let t = interpolation_factor();
79    let position_t = vec3<f32>(t);
80    let rotation_t = vec4<f32>(t);
81
82    let lhs_position = get_position(index);
83    let rhs_position = get_rhs_position(index);
84    let lhs_visibility = get_visibility(index);
85    let rhs_visibility = get_rhs_visibility(index);
86
87    let position = mix(lhs_position, rhs_position, position_t);
88    let visibility = mix(lhs_visibility, rhs_visibility, t);
89    set_output_position_visibility(index, position, visibility);
90
91    var sh = get_spherical_harmonics(index);
92    let rhs_sh = get_rhs_spherical_harmonics(index);
93    for (var i = 0u; i < #{SH_COEFF_COUNT}; i = i + 1u) {
94        sh[i] = mix(sh[i], rhs_sh[i], t);
95    }
96    set_output_spherical_harmonics(index, sh);
97
98#ifdef PRECOMPUTE_COVARIANCE_3D
99    var cov = get_cov3d(index);
100    let rhs_cov = get_rhs_cov3d(index);
101    for (var i = 0u; i < 6u; i = i + 1u) {
102        cov[i] = mix(cov[i], rhs_cov[i], t);
103    }
104    let opacity = mix(get_opacity(index), get_rhs_opacity(index), t);
105    set_output_covariance(index, cov, opacity);
106#else
107    let rotation = normalize_quaternion(
108        mix(
109            get_rotation(index),
110            get_rhs_rotation(index),
111            rotation_t,
112        ),
113    );
114
115    let scale = mix(get_scale(index), get_rhs_scale(index), position_t);
116    let opacity = mix(get_opacity(index), get_rhs_opacity(index), t);
117    set_output_transform(index, rotation, scale, opacity);
118#endif
119}