correlate.py

   1 #!/usr/bin/env python3
   2
   3 import gamma
   4 import numpy
   5 import perspective
   6 import scipy
   7 import scipy.ndimage
   8 import scipy.signal
   9 import sys
  10
  11 # a BLOCKS x BLOCKS section is correlated with a (BLOCKS + 1) x (BLOCKS + 1)
  12 # the maximum allowable shift is BLOCK_SIZE with BLOCKS of surrounding context
  13 BLOCKS = 4
  14 BLOCK_SIZE = 64
  15
  16 CUTOFF0 = 64
  17 CUTOFF1 = 4
  18
  19 in_jpg0 = 'tank_battle/down_2364.jpg'
  20 in_jpg1 = 'tank_battle/down_2365.jpg'
  21 out_jpg0 = 'out0_2364.jpg'
  22 out_jpg1 = 'out0_2365.jpg'
  23
  24 print(f'read {in_jpg0:s}')
  25 image0 = gamma.read_image(in_jpg0)
  26 shape = image0.shape
  27
  28 sys.stderr.write(f'write {out_jpg0:s}\n')
  29 gamma.write_image(out_jpg0, image0)
  30
  31 print(f'read {in_jpg1:s}')
  32 image1 = gamma.read_image(in_jpg1)
  33 assert image1.shape == shape
  34
  35 ys, xs, cs = shape
  36 xb = xs // BLOCK_SIZE
  37 yb = ys // BLOCK_SIZE
  38 print('xb', xb, 'yb', yb)
  39
  40 def bandpass(image):
  41   _, _, cs = image.shape
  42   return numpy.stack(
  43     [
  44       scipy.ndimage.gaussian_filter(image[:, :, i], CUTOFF1, mode = 'mirror') -
  45         scipy.ndimage.gaussian_filter(image[:, :, i], CUTOFF0, mode = 'mirror')
  46       for i in range(cs)
  47     ],
  48     2
  49   )
  50
  51 def correlate(image0_bp, x0, y0, image1_bp, x1, y1):
  52   block0 = image0_bp[
  53     y0:y0 + BLOCK_SIZE * BLOCKS,
  54     x0:x0 + BLOCK_SIZE * BLOCKS,
  55     :
  56   ]
  57   block1 = image1_bp[
  58     y1:y1 + BLOCK_SIZE * (BLOCKS + 1),
  59     x1:x1 + BLOCK_SIZE * (BLOCKS + 1),
  60     :
  61   ]
  62
  63   # note: swapping block1, block0 flips the output (subtracts x and y
  64   # from BLOCK_SIZE) and we need this to find matching part of image1
  65   corr = numpy.sum(
  66     numpy.stack(
  67       [
  68         scipy.signal.correlate(
  69           block1[:, :, i],
  70           block0[:, :, i],
  71           mode = 'valid'
  72         )
  73         for i in range(cs)
  74       ],
  75       0
  76     ),
  77     0
  78   )
  79   #temp = corr - numpy.mean(corr)
  80   #temp /= 10. * numpy.sqrt(numpy.mean(numpy.square(temp)))
  81   #gamma.write_image(f'corr_{i:d}_{j:d}.jpg', temp + .5)
  82
  83   y, x = numpy.unravel_index(numpy.argmax(corr), corr.shape)
  84   return (
  85     (x, y)
  86   if (
  87     x >= CUTOFF1 and
  88       x <= BLOCK_SIZE - CUTOFF1 and
  89       y >= CUTOFF1 and
  90       y <= BLOCK_SIZE - CUTOFF1
  91   ) else
  92     None
  93   )
  94
  95 print('bp0')
  96 image0_bp = bandpass(image0)
  97 gamma.write_image('image0_bp.jpg', image0_bp + .5)
  98
  99 print('bp1')
 100 image1_bp = bandpass(image1)
 101 gamma.write_image('image1_bp.jpg', image1_bp + .5)
 102
 103 print('align')
 104 buckets = [[[], []], [[], []]]
 105 for i in range(yb - BLOCKS):
 106   y1 = i * BLOCK_SIZE
 107   y0 = y1 + BLOCK_SIZE // 2
 108   for j in range(xb - BLOCKS):
 109     x1 = j * BLOCK_SIZE
 110     x0 = x1 + BLOCK_SIZE // 2
 111
 112     offset = correlate(image0_bp, x0, y0, image1_bp, x1, y1)
 113     if offset is not None:
 114       x, y = offset
 115       print('i', i, 'j', j, 'x', x, 'y', y)
 116       buckets[i >= (yb - BLOCKS) // 2][j >= (xb - BLOCKS) // 2].append(
 117         (x0, y0, x1, y1, x, y)
 118       )
 119
 120 p = []
 121 q = []
 122 for i in range(2):
 123   for j in range(2):
 124     k = len(buckets[i][j]) // 2
 125     xm = sorted([x for _, _, _, _, x, _ in buckets[i][j]])[k]
 126     ym = sorted([y for _, _, _, _, _, y in buckets[i][j]])[k]
 127     #print('i', i, 'j', j, 'xm', xm, 'ym', ym)
 128     u = numpy.array(
 129       [[x, y] for _, _, _, _, x, y in buckets[i][j]],
 130       numpy.double
 131     )
 132     v = numpy.array([xm, ym], numpy.double)
 133     k = numpy.argmin(numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1))
 134     x0, y0, x1, y1, x, y = buckets[i][j][k]
 135     #print('i', i, 'j', j, 'x', x, 'y', y)
 136     p.append(
 137       numpy.array(
 138         [
 139           x0 + (BLOCKS * BLOCK_SIZE // 2),
 140           y0 + (BLOCKS * BLOCK_SIZE // 2)
 141         ],
 142         numpy.double
 143       )
 144     )
 145     q.append(
 146       numpy.array(
 147         [
 148           x1 + x + (BLOCKS * BLOCK_SIZE // 2),
 149           y1 + y + (BLOCKS * BLOCK_SIZE // 2)
 150         ],
 151         numpy.double
 152       )
 153     )
 154 p = numpy.stack(p, 1)
 155 q = numpy.stack(q, 1)
 156 A = perspective.calc_transform(p, q)
 157
 158 print('remap')
 159 out_image1 = perspective.remap_image(A, image1)
 160
 161 sys.stderr.write(f'write {out_jpg1:s}\n')
 162 gamma.write_image(out_jpg1, out_image1)