Subroutinize the main parts of the auto-alignment routine

diff --git a/correlate.py b/correlate.py
index e56c77f..fa9abee 100755 (executable)
--- a/correlate.py
+++ b/correlate.py
@@ -37,40 +37,67 @@ xb = xs // BLOCK_SIZE
 yb = ys // BLOCK_SIZE
 print('xb', xb, 'yb', yb)
 
+def bandpass(image):
+  _, _, cs = image.shape
+  return numpy.stack(
+    [
+      scipy.ndimage.gaussian_filter(image[:, :, i], CUTOFF1, mode = 'mirror') -
+        scipy.ndimage.gaussian_filter(image[:, :, i], CUTOFF0, mode = 'mirror')
+      for i in range(cs)
+    ],
+    2
+  )
+
+def correlate(image0_bp, x0, y0, image1_bp, x1, y1):
+  block0 = image0_bp[
+    y0:y0 + BLOCK_SIZE * BLOCKS,
+    x0:x0 + BLOCK_SIZE * BLOCKS,
+    :
+  ]
+  block1 = image1_bp[
+    y1:y1 + BLOCK_SIZE * (BLOCKS + 1),
+    x1:x1 + BLOCK_SIZE * (BLOCKS + 1),
+    :
+  ]
+
+  # note: swapping block1, block0 flips the output (subtracts x and y
+  # from BLOCK_SIZE) and we need this to find matching part of image1
+  corr = numpy.sum(
+    numpy.stack(
+      [
+        scipy.signal.correlate(
+          block1[:, :, i],
+          block0[:, :, i],
+          mode = 'valid'
+        )
+        for i in range(cs)
+      ],
+      0
+    ),
+    0
+  )
+  #temp = corr - numpy.mean(corr)
+  #temp /= 10. * numpy.sqrt(numpy.mean(numpy.square(temp)))
+  #gamma.write_image(f'corr_{i:d}_{j:d}.jpg', temp + .5)
+
+  y, x = numpy.unravel_index(numpy.argmax(corr), corr.shape)
+  return (
+    (x, y)
+  if (
+    x >= CUTOFF1 and
+      x <= BLOCK_SIZE - CUTOFF1 and
+      y >= CUTOFF1 and
+      y <= BLOCK_SIZE - CUTOFF1
+  ) else
+    None
+  )
+
 print('bp0')
-image0_lp0 = numpy.stack(
-  [
-    scipy.ndimage.gaussian_filter(image0[:, :, i], CUTOFF0, mode = 'mirror')
-    for i in range(cs)
-  ],
-  2
-)
-image0_lp1 = numpy.stack(
-  [
-    scipy.ndimage.gaussian_filter(image0[:, :, i], CUTOFF1, mode = 'mirror')
-    for i in range(cs)
-  ],
-  2
-)
-image0_bp = image0_lp1 - image0_lp0
+image0_bp = bandpass(image0)
 gamma.write_image('image0_bp.jpg', image0_bp + .5)
 
 print('bp1')
-image1_lp0 = numpy.stack(
-  [
-    scipy.ndimage.gaussian_filter(image1[:, :, i], CUTOFF0, mode = 'mirror')
-    for i in range(cs)
-  ],
-  2
-)
-image1_lp1 = numpy.stack(
-  [
-    scipy.ndimage.gaussian_filter(image1[:, :, i], CUTOFF1, mode = 'mirror')
-    for i in range(cs)
-  ],
-  2
-)
-image1_bp = image1_lp1 - image1_lp0
+image1_bp = bandpass(image1)
 gamma.write_image('image1_bp.jpg', image1_bp + .5)
 
 print('align')
@@ -82,44 +109,10 @@ for i in range(yb - BLOCKS):
     x1 = j * BLOCK_SIZE
     x0 = x1 + BLOCK_SIZE // 2
 
-    block0 = image0_bp[
-      y0:y0 + BLOCK_SIZE * BLOCKS,
-      x0:x0 + BLOCK_SIZE * BLOCKS,
-      :
-    ]
-    block1 = image1_bp[
-      y1:y1 + BLOCK_SIZE * (BLOCKS + 1),
-      x1:x1 + BLOCK_SIZE * (BLOCKS + 1),
-      :
-    ]
-
-    # note: swapping block1, block0 flips the output (subtracts x and y
-    # from BLOCK_SIZE) and we need this to find matching part of image1
-    corr = numpy.sum(
-      numpy.stack(
-        [
-          scipy.signal.correlate(
-            block1[:, :, i],
-            block0[:, :, i],
-            mode = 'valid'
-          )
-          for i in range(cs)
-        ],
-        0
-      ),
-      0
-    )
-    #temp = corr - numpy.mean(corr)
-    #temp /= 10. * numpy.sqrt(numpy.mean(numpy.square(temp)))
-    #gamma.write_image(f'corr_{i:d}_{j:d}.jpg', temp + .5)
-
-    y, x = numpy.unravel_index(numpy.argmax(corr), corr.shape)
-    if (
-      y >= CUTOFF1 and
-        y <= BLOCK_SIZE - CUTOFF1 and
-        x >= CUTOFF1 and
-        x <= BLOCK_SIZE - CUTOFF1
-    ):
+    offset = correlate(image0_bp, x0, y0, image1_bp, x1, y1)
+    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(
         (x0, y0, x1, y1, x, y)
       )
@@ -163,20 +156,7 @@ q = numpy.stack(q, 1)
 A = perspective.calc_transform(p, q)
 
 print('remap')
-coords = numpy.zeros((2, ys, xs), numpy.double)
-coords[0, :, :] = numpy.arange(ys, dtype = numpy.double)[:, numpy.newaxis]
-coords[1, :, :] = numpy.arange(xs, dtype = numpy.double)[numpy.newaxis, :]
-mapped_coords = perspective.apply_transform_multi(
-  A,
-  coords[::-1, :, :].reshape((2, ys * xs))
-).reshape(2, ys, xs)[::-1, :, :]
-out_image1 = numpy.stack(
-  [
-    scipy.ndimage.map_coordinates(image1[:, :, j], mapped_coords)
-    for j in range(cs)
-  ],
-  2
-)
+out_image1 = perspective.remap_image(A, image1)
 
 sys.stderr.write(f'write {out_jpg1:s}\n')
 gamma.write_image(out_jpg1, out_image1)

diff --git a/perspective.py b/perspective.py
index 9128fb7..951d960 100755 (executable)
--- a/perspective.py
+++ b/perspective.py
@@ -2,6 +2,8 @@
 
 import numpy
 import numpy.linalg
+import scipy
+import scipy.ndimage
 
 # call with 2x4 vector (4 points of x, y)
 def calc_basis(p):
@@ -47,6 +49,23 @@ def apply_transform_multi(A, p):
 def apply_transform(A, p):
   return apply_transform_multi(A, p[:, numpy.newaxis])[:, 0]
 
+def remap_image(A, image):
+  ys, xs, cs = image.shape
+  coords = numpy.zeros((2, ys, xs), numpy.double)
+  coords[0, :, :] = numpy.arange(ys, dtype = numpy.double)[:, numpy.newaxis]
+  coords[1, :, :] = numpy.arange(xs, dtype = numpy.double)[numpy.newaxis, :]
+  mapped_coords = apply_transform_multi(
+    A,
+    coords[::-1, :, :].reshape((2, ys * xs))
+  ).reshape(2, ys, xs)[::-1, :, :]
+  return numpy.stack(
+    [
+      scipy.ndimage.map_coordinates(image[:, :, j], mapped_coords)
+      for j in range(cs)
+    ],
+    2
+  )
+
 if __name__ == '__main__':
   x = numpy.array([[11., 52., 23., 74.], [15., 36., 27., 58.]], numpy.double)
   y = numpy.array([[31., 92., 73., 24.], [65., 26., 87., 18.]], numpy.double)