Combine the correlation loop with the trend-finding loop (no need for buckets)

diff --git a/correlate.py b/correlate.py
index 1e93a9f..da13367 100755 (executable)
--- a/correlate.py
+++ b/correlate.py
@@ -33,8 +33,8 @@ image1 = gamma.read_image(in_jpg1)
 assert image1.shape == shape
 
 ys, xs, cs = shape
-xb = xs // BLOCK_SIZE
-yb = ys // BLOCK_SIZE
+xb = (xs // BLOCK_SIZE - BLOCKS) // 2
+yb = (ys // BLOCK_SIZE - BLOCKS) // 2
 print('xb', xb, 'yb', yb)
 
 def bandpass(image):
@@ -108,30 +108,43 @@ gamma.write_image('image1_bp.jpg', image1_bp + .5)
 print('align')
 buckets = [[[], []], [[], []]]
 for i in range(yb - BLOCKS):
-  yc = i * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
   for j in range(xb - BLOCKS):
     xc = j * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
-    offset = correlate(image0_bp, image1_bp, xc, yc)
-    if offset is not None:
-      x, y = offset
-      print('i', i, 'j', j, 'x', x, 'y', y)
-      buckets[i >= (yb - BLOCKS) // 2][j >= (xb - BLOCKS) // 2].append(
-        (xc, yc, xc + x, yc + y)
-      )
 
 p = []
 q = []
 for i in range(2):
   for j in range(2):
-    k = len(buckets[i][j]) // 2
-    offset = [(xc1 - xc0, yc1 - yc0) for xc0, yc0, xc1, yc1 in buckets[i][j]]
+    # correlate blocks in (i, j)-corner
+    blocks = []
+    for k in range(yb):
+      yc = k * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
+      if i:
+        yc = ys - yc
+      for l in range(xb):
+        xc = l * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
+        if j:
+          xc = xs - xc
+        offset = correlate(image0_bp, image1_bp, xc, yc)
+        if offset is not None:
+          x, y = offset
+          print('i', i, 'j', j, 'k', k, 'l', l, 'x', x, 'y', y)
+          blocks.append((xc, yc, xc + x, yc + y))
+
+    # 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]
     #print('i', i, 'j', j, 'xm', xm, 'ym', ym)
+
+    # choose block closest to trend in (i, j)-corner
     u = numpy.array(offset, numpy.double)
     v = numpy.array([xm, ym], numpy.double)
     k = numpy.argmin(numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1))
-    xc0, yc0, xc1, yc1 = buckets[i][j][k]
+
+    # 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))