1 files changed, 47 insertions, 6 deletions

diff --git a/flowy/src/motion/enc.wgsl b/flowy/src/motion/enc.wgsl
index bb0139f..b94bac9 100644
--- a/flowy/src/motion/enc.wgsl
+++ b/flowy/src/motion/enc.wgsl
@@ -1,12 +1,15 @@
 
 struct Params {
     block_size: vec2<i32>,
-    offsets_stride: u32
+    offsets_stride: u32,
+    search_radius: i32,
+    skip_threshold: f32,
 }
 
 struct BlockOffset {
     score: f32,
     offset: vec2<i32>,
+    tint: vec3<f32>,
 }
 
 @group(0) @binding(0) var<uniform> params: Params;
@@ -17,10 +20,46 @@ struct BlockOffset {
 @compute @workgroup_size(1) fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
     let uv = vec2<i32>(global_id.xy) * params.block_size;
 
+    let _f = search_offset(uv);
+    let best_err = _f.error;
+    let best_offset = _f.offset;
+
+    var best_tint = vec3(0.);
+
+    var average_pcol = vec3(0.);
+    var average_ncol = vec3(0.);
+    for (var x = 0; x < params.block_size.x; x++) {
+    for (var y = 0; y < params.block_size.y; y++) {
+        let base = uv+vec2(x,y);
+        average_pcol += textureLoad(prev, base+best_offset, 0).rgb;
+        average_ncol += textureLoad(next, base, 0).rgb;
+    }}
+
+    let tint = (average_ncol - average_pcol) / f32(params.block_size.x * params.block_size.y);
+
+    var err = 0.;
+    for (var x = 0; x < params.block_size.x; x++) {
+    for (var y = 0; y < params.block_size.y; y++) {
+        let base = uv+vec2(x,y);
+        let pcol = textureLoad(prev, base+best_offset, 0).rgb+tint;
+        let ncol = textureLoad(next, base, 0).rgb;
+        err += distance(pcol, ncol);
+    }}
+    if err < best_err {
+        best_tint = tint;
+    }
+
+    offsets[global_id.x + global_id.y * params.offsets_stride] = BlockOffset(best_err, best_offset, best_tint);
+}
+
+struct SearchRes {offset: vec2<i32>, error: f32}
+
+fn search_offset(uv: vec2<i32>) -> SearchRes {
     var best_err = 100000000.;
-    var best_offset = vec2(0,0);
-    for (var ox = -5; ox <= 5; ox++) {
-    for (var oy = -5; oy <= 5; oy++) {
+    var best_offset = vec2(0);
+    // TODO: better ordering
+    for (var ox = -params.search_radius; ox <= params.search_radius; ox++) {
+    for (var oy = -params.search_radius; oy <= params.search_radius; oy++) {
         let offset = vec2(ox,oy);
 
         var err = 0.;
@@ -35,10 +74,12 @@ struct BlockOffset {
         if err < best_err {
             best_err = err;
             best_offset = offset;
+            if err < params.skip_threshold {
+                return SearchRes(offset, err);
+            }
         }
     }}
-
-    offsets[global_id.x + global_id.y * params.offsets_stride] = BlockOffset(best_err, best_offset);
+    return SearchRes(best_offset,best_err);
 }
 
 // fn colormap_vec(v: vec2<f32>) -> vec3<f32> {