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