X-Git-Url: https://git.sesse.net/?a=blobdiff_plain;f=motion_search.frag;h=05e901c12844cfa3dd38753435c4b8e649fd0144;hb=670611b5a707fb8dcebf60fcfcd0930545d14875;hp=e31a673484be9479f71a7428feec4348604bbf5d;hpb=bcdff7b1a183d7ff138a920997182d00f2decf0e;p=nageru

diff --git a/motion_search.frag b/motion_search.frag
index e31a673..05e901c 100644
--- a/motion_search.frag
+++ b/motion_search.frag
@@ -36,21 +36,24 @@
  */
 
 const uint patch_size = 12;
-const uint num_iterations = 16;
+const uint num_iterations = 8;
 
 in vec2 flow_tc;
-in vec2 patch_bottom_left_texel;  // Center of bottom-left texel of patch.
-out vec2 out_flow;
+in vec2 patch_center;
+out vec3 out_flow;
 
 uniform sampler2D flow_tex, grad0_tex, image0_tex, image1_tex;
-uniform float image_width, image_height, inv_image_width, inv_image_height;
+uniform vec2 inv_image_size, inv_prev_level_size;
 
 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 * vec2(image_width, image_height))
-		* vec2(inv_image_width, inv_image_height);
+	vec2 image_size = textureSize(image0_tex, 0);
+
+	// 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.
 	// This is the part where we really save by the inverse search
@@ -63,15 +66,18 @@ 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;
-			tc.x = base.x + x * inv_image_width;
-			tc.y = base.y + y * inv_image_height;
+			vec2 tc = base + uvec2(x, y) * inv_image_size;
 			vec2 grad = texture(grad0_tex, tc).xy;
 			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];
@@ -87,38 +93,67 @@ 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;
-	}
+	// 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;
-				tc.x = base.x + x * inv_image_width;
-				tc.y = base.y + y * inv_image_height;
+				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;
+				float warped = texture(image1_tex, tc + u_norm).x;
 				du += grad * (warped - t);
+				warped_sum += warped;
 			}
 		}
-		u += (H_inv * du) * vec2(inv_image_width, inv_image_height);
+
+		// 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.
+		mean_diff = (warped_sum - template_sum) * (1.0 / (patch_size * patch_size));
+		du -= grad_sum * mean_diff;
+
+		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) * vec2(image_width, image_height)) > 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);
 }