Proof of concept automatic alignment routine using correlation

[stop_motion.git] / correlate.py

diff --git a/correlate.py b/correlate.py
new file mode 100755 (executable)
index 0000000..e56c77f
--- /dev/null
+++ b/correlate.py
@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+
+import gamma
+import numpy
+import perspective
+import scipy
+import scipy.ndimage
+import scipy.signal
+import sys
+
+# a BLOCKS x BLOCKS section is correlated with a (BLOCKS + 1) x (BLOCKS + 1)
+# the maximum allowable shift is BLOCK_SIZE with BLOCKS of surrounding context
+BLOCKS = 4
+BLOCK_SIZE = 64
+
+CUTOFF0 = 64
+CUTOFF1 = 4
+
+in_jpg0 = 'tank_battle/down_2364.jpg'
+in_jpg1 = 'tank_battle/down_2365.jpg'
+out_jpg0 = 'out0_2364.jpg'
+out_jpg1 = 'out0_2365.jpg'
+
+print(f'read {in_jpg0:s}')
+image0 = gamma.read_image(in_jpg0)
+shape = image0.shape
+
+sys.stderr.write(f'write {out_jpg0:s}\n')
+gamma.write_image(out_jpg0, image0)
+
+print(f'read {in_jpg1:s}')
+image1 = gamma.read_image(in_jpg1)
+assert image1.shape == shape
+
+ys, xs, cs = shape
+xb = xs // BLOCK_SIZE
+yb = ys // BLOCK_SIZE
+print('xb', xb, 'yb', yb)
+
+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
+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
+gamma.write_image('image1_bp.jpg', image1_bp + .5)
+
+print('align')
+buckets = [[[], []], [[], []]]
+for i in range(yb - BLOCKS):
+  y1 = i * BLOCK_SIZE
+  y0 = y1 + BLOCK_SIZE // 2
+  for j in range(xb - 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
+    ):
+      buckets[i >= (yb - BLOCKS) // 2][j >= (xb - BLOCKS) // 2].append(
+        (x0, y0, x1, y1, x, y)
+      )
+
+p = []
+q = []
+for i in range(2):
+  for j in range(2):
+    k = len(buckets[i][j]) // 2
+    xm = sorted([x for _, _, _, _, x, _ in buckets[i][j]])[k]
+    ym = sorted([y for _, _, _, _, _, y in buckets[i][j]])[k]
+    #print('i', i, 'j', j, 'xm', xm, 'ym', ym)
+    u = numpy.array(
+      [[x, y] for _, _, _, _, x, y in buckets[i][j]],
+      numpy.double
+    )
+    v = numpy.array([xm, ym], numpy.double)
+    k = numpy.argmin(numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1))
+    x0, y0, x1, y1, x, y = buckets[i][j][k]
+    #print('i', i, 'j', j, 'x', x, 'y', y)
+    p.append(
+      numpy.array(
+        [
+          x0 + (BLOCKS * BLOCK_SIZE // 2),
+          y0 + (BLOCKS * BLOCK_SIZE // 2)
+        ],
+        numpy.double
+      )
+    )
+    q.append(
+      numpy.array(
+        [
+          x1 + x + (BLOCKS * BLOCK_SIZE // 2),
+          y1 + y + (BLOCKS * BLOCK_SIZE // 2)
+        ],
+        numpy.double
+      )
+    )
+p = numpy.stack(p, 1)
+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
+)
+
+sys.stderr.write(f'write {out_jpg1:s}\n')
+gamma.write_image(out_jpg1, out_image1)