1#define_import_path bevy_gaussian_splatting::optical_flow
2
3#import bevy_pbr::{
4    forward_io::VertexOutput,
5    prepass_utils,
6}
7#import bevy_render::color_operations::hsv_to_rgb
8#import bevy_render::maths::PI_2
9
10#import bevy_gaussian_splatting::bindings::{
11    globals,
12    previous_view_uniforms,
13    view,
14}
15
16
17
18fn calculate_motion_vector(
19    world_position: vec3<f32>,
20    previous_world_position: vec3<f32>,
21) -> vec2<f32> {
22    let world_position_t = vec4<f32>(world_position, 1.0);
23    let previous_world_position_t = vec4<f32>(previous_world_position, 1.0);
24    let clip_position_t = view.unjittered_clip_from_world * world_position_t;
25    let clip_position = clip_position_t.xy / clip_position_t.w;
26    let previous_clip_position_t = previous_view_uniforms.clip_from_world * previous_world_position_t;
27    let previous_clip_position = previous_clip_position_t.xy / previous_clip_position_t.w;
28    // These motion vectors are used as offsets to UV positions and are stored
29    // in the range -1,1 to allow offsetting from the one corner to the
30    // diagonally-opposite corner in UV coordinates, in either direction.
31    // A difference between diagonally-opposite corners of clip space is in the
32    // range -2,2, so this needs to be scaled by 0.5. And the V direction goes
33    // down where clip space y goes up, so y needs to be flipped.
34    return (clip_position - previous_clip_position) * vec2(0.5, -0.5);
35}
36
37
38fn optical_flow_to_rgb(
39    motion_vector: vec2<f32>,
40) -> vec3<f32> {
41    let flow = motion_vector / globals.delta_time;
42
43    let radius = length(flow);
44    var angle = atan2(flow.y, flow.x);
45    if (angle < 0.0) {
46        angle += PI_2;
47    }
48
49    // let sigma: f32 = 0.15;
50    // let norm_factor = sigma * 2.0;
51    // let m = clamp(radius / norm_factor, 0.0, 1.0);
52    let m = clamp(radius, 0.0, 1.0);
53
54    let rgb = hsv_to_rgb(vec3<f32>(angle, m, 1.0));
55    return rgb;
56}
57
58// TODO: support immediate vs. persistent previous_view, aiding with no-smoothness on the pan-orbit camera (required by cpu sort)
59// TODO: set clear color to white in optical flow render mode