]> git.sesse.net Git - nageru/blobdiff - motion_search.frag
Allow symlinked frame files. Useful for testing.
[nageru] / motion_search.frag
index d9d1f4eb4997c7650fb0f4902acf2c42b5ff4e4c..eb4f7c77fe2c6f8cb056a46caeaa89d58f0ee486 100644 (file)
   be ideal.
  */
 
-const uint patch_size = 12;
-const uint num_iterations = 16;
+in vec3 flow_tc;
+in vec2 patch_center;
+flat in int ref_layer, search_layer;
+out vec3 out_flow;
+
+uniform sampler2DArray flow_tex, image_tex;
+uniform usampler2DArray grad_tex;  // Also contains the corresponding reference image.
+uniform vec2 inv_image_size, inv_prev_level_size;
+uniform uint patch_size;
+uniform uint num_iterations;
+
+vec3 unpack_gradients(uint v)
+{
+       uint vi = v & 0xffu;
+       uint xi = (v >> 8) & 0xfffu;
+       uint yi = v >> 20;
+       vec3 r = vec3(xi * (1.0f / 4095.0f) - 0.5f, yi * (1.0f / 4095.0f) - 0.5f, vi * (1.0f / 255.0f));
+       return r;
+}
+
+// Note: The third variable is the actual pixel value.
+vec3 get_gradients(vec3 tc)
+{
+       vec3 grad = unpack_gradients(texture(grad_tex, tc).x);
 
-in vec2 flow_tc;
-in vec2 patch_bottom_left_texel;  // Center of bottom-left texel of patch.
-out vec2 out_flow;
+       // Zero gradients outside the image. (We'd do this with a sampler,
+       // but we want the repeat behavior for the actual texels, in the
+       // z channel.)
+       if (any(lessThan(tc.xy, vec2(0.0f))) || any(greaterThan(tc.xy, vec2(1.0f)))) {
+               grad.xy = vec2(0.0f);
+       }
 
-uniform sampler2D flow_tex, grad0_tex, image0_tex, image1_tex;
-uniform vec2 image_size, inv_image_size;
+       return grad;
+}
 
 void main()
 {
-       // Lock patch_bottom_left_texel to an integer, so that we never get
-       // any bilinear artifacts for the gradient.
-       vec2 base = round(patch_bottom_left_texel * image_size)
+       vec2 image_size = textureSize(grad_tex, 0).xy;
+
+       // Lock the patch center to an integer, so that we never get
+       // any bilinear artifacts for the gradient. (NOTE: This assumes an
+       // even patch size.) Then calculate the bottom-left texel of the patch.
+       vec2 base = (round(patch_center * image_size) - (0.5f * patch_size - 0.5f))
                * inv_image_size;
 
        // First, precompute the pseudo-Hessian for the template patch.
@@ -68,13 +96,13 @@ void main()
        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;
-                       vec2 grad = texture(grad0_tex, tc).xy;
+                       vec3 grad = get_gradients(vec3(tc, ref_layer));
                        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;
+                       template_sum += grad.z;  // The actual template pixel value.
+                       grad_sum += grad.xy;
                }
        }
        H[1][0] = H[0][1];
@@ -90,32 +118,22 @@ void main()
 
        mat2 H_inv = inverse(H);
 
-       // Fetch the initial guess for the flow. (We need the normalization step
-       // because densification works by accumulating; see the comments on the
-       // Densify class.)
-       vec3 prev_flow = texture(flow_tex, flow_tc).xyz;
-       vec2 initial_u;
-       if (prev_flow.z < 1e-3) {
-               initial_u = vec2(0.0, 0.0);
-       } 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.
+       // Fetch the initial guess for the flow, and convert from the previous size to this one.
+       vec2 initial_u = texture(flow_tex, flow_tc).xy * (image_size * inv_prev_level_size);
        vec2 u = initial_u;
+       float mean_diff, first_mean_diff;
 
        for (uint i = 0; i < num_iterations; ++i) {
                vec2 du = vec2(0.0, 0.0);
                float warped_sum = 0.0f;
+               vec2 u_norm = u * inv_image_size;  // In [0..1] coordinates instead of pixels.
                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;
-                               vec2 grad = texture(grad0_tex, tc).xy;
-                               float t = texture(image0_tex, tc).x;
-                               float warped = texture(image1_tex, tc + u).x;
-                               du += grad * (warped - t);
+                               vec3 grad = get_gradients(vec3(tc, ref_layer));
+                               float t = grad.z;
+                               float warped = texture(image_tex, vec3(tc + u_norm, search_layer)).x;
+                               du += grad.xy * (warped - t);
                                warped_sum += warped;
                        }
                }
@@ -131,15 +149,36 @@ void main()
                //   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));
+               mean_diff = (warped_sum - template_sum) * (1.0 / float(patch_size * patch_size));
+               du -= grad_sum * mean_diff;
 
-               u += (H_inv * du) * inv_image_size;
+               if (i == 0) {
+                       first_mean_diff = mean_diff;
+               }
+
+               // Do the actual update.
+               u -= H_inv * du;
        }
 
-       // Reject if we moved too far.
-       if (length((u - initial_u) * image_size) > patch_size) {
+       // Reject if we moved too far. Note that the paper says “too far” is the
+       // patch size, but the DIS code uses half of a patch size. The latter seems
+       // to give much better overall results.
+       //
+       // Also reject if the patch goes out-of-bounds (the paper does not mention this,
+       // but the code does, and it seems to be critical to avoid really bad behavior
+       // at the edges).
+       vec2 patch_center = (base * image_size - 0.5f) + patch_size * 0.5f + u;
+       if (length(u - initial_u) > (patch_size * 0.5f) ||
+           patch_center.x < -(patch_size * 0.5f) ||
+           image_size.x - patch_center.x < -(patch_size * 0.5f) ||
+           patch_center.y < -(patch_size * 0.5f) ||
+           image_size.y - patch_center.y < -(patch_size * 0.5f)) {
                u = initial_u;
+               mean_diff = first_mean_diff;
        }
 
-       out_flow = u;
+       // NOTE: The mean patch diff will be for the second-to-last patch,
+       // not the true position of du. But hopefully, it will be very close.
+       u *= inv_image_size;
+       out_flow = vec3(u.x, u.y, mean_diff);
 }