Add mean normalization for patches.

[nageru] / motion_search.frag

diff --git a/motion_search.frag b/motion_search.frag
index 5f2b6f8c2e1c79ffaee0de5b9c04ace4153b0f1d..d9d1f4eb4997c7650fb0f4902acf2c42b5ff4e4c 100644 (file)
--- a/motion_search.frag
+++ b/motion_search.frag
@@ -63,6 +63,8 @@ void main()
        // this is an outer product, so we get a (symmetric) 2x2 matrix,
        // not a scalar.
        mat2 H = mat2(0.0f);
+       vec2 grad_sum = vec2(0.0f);  // Used for patch normalization.
+       float template_sum = 0.0f;
        for (uint y = 0; y < patch_size; ++y) {
                for (uint x = 0; x < patch_size; ++x) {
                        vec2 tc = base + uvec2(x, y) * inv_image_size;
@@ -70,6 +72,9 @@ void main()
                        H[0][0] += grad.x * grad.x;
                        H[1][1] += grad.y * grad.y;
                        H[0][1] += grad.x * grad.y;
+
+                       template_sum += texture(image0_tex, tc).x;
+                       grad_sum += grad;
                }
        }
        H[1][0] = H[0][1];
@@ -95,10 +100,15 @@ void main()
        } else {
                initial_u = prev_flow.xy / prev_flow.z;
        }
+
+       // Note: The flow is in OpenGL coordinates [0..1], but the calculations
+       // generally come out in pixels since the gradient is in pixels,
+       // so we need to convert at the end.
        vec2 u = initial_u;
 
        for (uint i = 0; i < num_iterations; ++i) {
                vec2 du = vec2(0.0, 0.0);
+               float warped_sum = 0.0f;
                for (uint y = 0; y < patch_size; ++y) {
                        for (uint x = 0; x < patch_size; ++x) {
                                vec2 tc = base + uvec2(x, y) * inv_image_size;
@@ -106,8 +116,23 @@ void main()
                                float t = texture(image0_tex, tc).x;
                                float warped = texture(image1_tex, tc + u).x;
                                du += grad * (warped - t);
+                               warped_sum += warped;
                        }
                }
+
+               // Subtract the mean for patch normalization. We've done our
+               // sums without subtracting the means (because we didn't know them
+               // beforehand), ie.:
+               //
+               //   sum(S^T * ((x + µ1) - (y + µ2))) = sum(S^T * (x - y)) + (µ1 – µ2) sum(S^T)
+               //
+               // which gives trivially
+               //
+               //   sum(S^T * (x - y)) = [what we calculated] - (µ1 - µ2) sum(S^T)
+               //
+               // so we can just subtract away the mean difference here.
+               du -= grad_sum * (warped_sum - template_sum) * (1.0 / (patch_size * patch_size));
+
                u += (H_inv * du) * inv_image_size;
        }