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 - BLOCKS) // 2
  37 yb = (ys // BLOCK_SIZE - BLOCKS) // 2
  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, image1_bp, xc, yc):
  52   x0 = xc - BLOCKS * BLOCK_SIZE // 2
  53   y0 = yc - BLOCKS * BLOCK_SIZE // 2
  54   x1 = xc - (BLOCKS + 1) * BLOCK_SIZE // 2
  55   y1 = yc - (BLOCKS + 1) * BLOCK_SIZE // 2
  56
  57   block0 = image0_bp[
  58     y0:y0 + BLOCK_SIZE * BLOCKS,
  59     x0:x0 + BLOCK_SIZE * BLOCKS,
  60     :
  61   ]
  62   block1 = image1_bp[
  63     y1:y1 + (BLOCKS + 1) * BLOCK_SIZE,
  64     x1:x1 + (BLOCKS + 1) * BLOCK_SIZE,
  65     :
  66   ]
  67
  68   # note: swapping block1, block0 flips the output (subtracts x and y
  69   # from BLOCK_SIZE) and we need this to find matching part of image1
  70   corr = numpy.sum(
  71     numpy.stack(
  72       [
  73         scipy.signal.correlate(
  74           block1[:, :, i],
  75           block0[:, :, i],
  76           mode = 'valid'
  77         )
  78         for i in range(cs)
  79       ],
  80       0
  81     ),
  82     0
  83   )
  84   #temp = corr - numpy.mean(corr)
  85   #temp /= 10. * numpy.sqrt(numpy.mean(numpy.square(temp)))
  86   #gamma.write_image(f'corr_{i:d}_{j:d}.jpg', temp + .5)
  87
  88   y, x = numpy.unravel_index(numpy.argmax(corr), corr.shape)
  89   return (
  90     (x - BLOCK_SIZE // 2, y - BLOCK_SIZE // 2)
  91   if (
  92     x >= CUTOFF1 and
  93       x <= BLOCK_SIZE - CUTOFF1 and
  94       y >= CUTOFF1 and
  95       y <= BLOCK_SIZE - CUTOFF1
  96   ) else
  97     None
  98   )
  99
 100 print('bp0')
 101 image0_bp = bandpass(image0)
 102 gamma.write_image('image0_bp.jpg', image0_bp + .5)
 103
 104 print('bp1')
 105 image1_bp = bandpass(image1)
 106 gamma.write_image('image1_bp.jpg', image1_bp + .5)
 107
 108 print('align')
 109 buckets = [[[], []], [[], []]]
 110 for i in range(yb - BLOCKS):
 111   for j in range(xb - BLOCKS):
 112     xc = j * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
 113
 114 p = []
 115 q = []
 116 for i in range(2):
 117   for j in range(2):
 118     # correlate blocks in (i, j)-corner
 119     blocks = []
 120     for k in range(yb):
 121       yc = k * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
 122       if i:
 123         yc = ys - yc
 124       for l in range(xb):
 125         xc = l * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
 126         if j:
 127           xc = xs - xc
 128         offset = correlate(image0_bp, image1_bp, xc, yc)
 129         if offset is not None:
 130           x, y = offset
 131           print('i', i, 'j', j, 'k', k, 'l', l, 'x', x, 'y', y)
 132           blocks.append((xc, yc, xc + x, yc + y))
 133
 134     # find the trend in (i, j)-corner
 135     k = len(blocks) // 2
 136     offset = [(xc1 - xc0, yc1 - yc0) for xc0, yc0, xc1, yc1 in blocks]
 137     xm = sorted([x for x, _ in offset])[k]
 138     ym = sorted([y for _, y in offset])[k]
 139     #print('i', i, 'j', j, 'xm', xm, 'ym', ym)
 140
 141     # choose block closest to trend in (i, j)-corner
 142     u = numpy.array(offset, numpy.double)
 143     v = numpy.array([xm, ym], numpy.double)
 144     k = numpy.argmin(numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1))
 145
 146     # output centres of chosen block for the perspective transform
 147     xc0, yc0, xc1, yc1 = blocks[k]
 148     #print('i', i, 'j', j, 'x', x, 'y', y)
 149     p.append(numpy.array([xc0, yc0], numpy.double))
 150     q.append(numpy.array([xc1, yc1], numpy.double))
 151 p = numpy.stack(p, 1)
 152 q = numpy.stack(q, 1)
 153 A = perspective.calc_transform(p, q)
 154
 155 print('remap')
 156 out_image1 = perspective.remap_image(A, image1)
 157
 158 sys.stderr.write(f'write {out_jpg1:s}\n')
 159 gamma.write_image(out_jpg1, out_image1)