Add corner candidates to fine-tune the perspective fit

diff --git a/correlate.py b/correlate.py
index da13367..d5f3aeb 100755 (executable)
--- a/correlate.py
+++ b/correlate.py
@@ -13,6 +13,8 @@ import sys
 BLOCKS = 4
 BLOCK_SIZE = 64
 
+CORNER_CANDIDATES = 8
+
 CUTOFF0 = 64
 CUTOFF1 = 4
 
@@ -105,17 +107,19 @@ print('bp1')
 image1_bp = bandpass(image1)
 gamma.write_image('image1_bp.jpg', image1_bp + .5)
 
-print('align')
+print('find corner candidates')
 buckets = [[[], []], [[], []]]
 for i in range(yb - BLOCKS):
   for j in range(xb - BLOCKS):
     xc = j * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
 
-p = []
-q = []
+p_all = []
+q_all = []
+corner_candidates = []
 for i in range(2):
   for j in range(2):
     # correlate blocks in (i, j)-corner
+    offsets = []
     blocks = []
     for k in range(yb):
       yc = k * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
@@ -127,33 +131,72 @@ for i in range(2):
           xc = xs - xc
         offset = correlate(image0_bp, image1_bp, xc, yc)
         if offset is not None:
+          offsets.append(offset)
           x, y = offset
-          print('i', i, 'j', j, 'k', k, 'l', l, 'x', x, 'y', y)
-          blocks.append((xc, yc, xc + x, yc + y))
+          xc1 = xc + x
+          yc1 = yc + y
+          #print('i', i, 'j', j, 'k', k, 'l', l, 'x', x, 'y', y)
+          p_all.append(numpy.array([xc, yc], numpy.double))
+          q_all.append(numpy.array([xc1, yc1], numpy.double))
+          blocks.append((xc, yc, xc1, yc1))
 
     # find the trend in (i, j)-corner
     k = len(blocks) // 2
-    offset = [(xc1 - xc0, yc1 - yc0) for xc0, yc0, xc1, yc1 in blocks]
-    xm = sorted([x for x, _ in offset])[k]
-    ym = sorted([y for _, y in offset])[k]
+    xm = sorted([x for x, _ in offsets])[k]
+    ym = sorted([y for _, y in offsets])[k]
     #print('i', i, 'j', j, 'xm', xm, 'ym', ym)
 
-    # choose block closest to trend in (i, j)-corner
-    u = numpy.array(offset, numpy.double)
+    # choose CORNER_CANDIDATES blocks closest to trend in (i, j)-corner
+    u = numpy.array(offsets, numpy.double)
     v = numpy.array([xm, ym], numpy.double)
-    k = numpy.argmin(numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1))
-
-    # output centres of chosen block for the perspective transform
-    xc0, yc0, xc1, yc1 = blocks[k]
-    #print('i', i, 'j', j, 'x', x, 'y', y)
-    p.append(numpy.array([xc0, yc0], numpy.double))
-    q.append(numpy.array([xc1, yc1], numpy.double))
-p = numpy.stack(p, 1)
-q = numpy.stack(q, 1)
-A = perspective.calc_transform(p, q)
+    dist = numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1)
+    corner_candidates.append(
+      sorted(
+        [(dist[i],) + blocks[i] for i in range(len(blocks))]
+      )[:CORNER_CANDIDATES]
+    )
+p_all = numpy.stack(p_all, 1)
+q_all = numpy.stack(q_all, 1)
+
+# try all combinations of the corner candidates
+print('try corner candidates')
+p = numpy.zeros((2, 4), numpy.double)
+q = numpy.zeros((2, 4), numpy.double)
+best_dist = None
+best_A = None
+for _, xc00, yc00, xc10, yc10 in corner_candidates[0]:
+  p[0, 0] = xc00
+  p[1, 0] = yc00
+  q[0, 0] = xc10
+  q[1, 0] = yc10
+  for _, xc01, yc01, xc11, yc11 in corner_candidates[1]:
+    p[0, 1] = xc01
+    p[1, 1] = yc01
+    q[0, 1] = xc11
+    q[1, 1] = yc11
+    for _, xc02, yc02, xc12, yc12 in corner_candidates[2]:
+      p[0, 2] = xc02
+      p[1, 2] = yc02
+      q[0, 2] = xc12
+      q[1, 2] = yc12
+      for _, xc03, yc03, xc13, yc13 in corner_candidates[3]:
+        p[0, 3] = xc03
+        p[1, 3] = yc03
+        q[0, 3] = xc13
+        q[1, 3] = yc13
+
+        A = perspective.calc_transform(p, q)
+        dist = numpy.sum(
+          numpy.square(
+            q_all - perspective.apply_transform_multi(A, p_all)
+          )
+        )
+        if best_dist is None or dist < best_dist:
+          best_dist = dist
+          best_A = A
 
 print('remap')
-out_image1 = perspective.remap_image(A, image1)
+out_image1 = perspective.remap_image(best_A, image1)
 
 sys.stderr.write(f'write {out_jpg1:s}\n')
 gamma.write_image(out_jpg1, out_image1)