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 CORNER_CANDIDATES = 8
  17
  18 CUTOFF0 = 64
  19 CUTOFF1 = 4
  20
  21 in_jpg0 = 'tank_battle/down_2364.jpg'
  22 in_jpg1 = 'tank_battle/down_2365.jpg'
  23 out_jpg0 = 'out0_2364.jpg'
  24 out_jpg1 = 'out0_2365.jpg'
  25
  26 print(f'read {in_jpg0:s}')
  27 image0 = gamma.read_image(in_jpg0)
  28 shape = image0.shape
  29
  30 sys.stderr.write(f'write {out_jpg0:s}\n')
  31 gamma.write_image(out_jpg0, image0)
  32
  33 print(f'read {in_jpg1:s}')
  34 image1 = gamma.read_image(in_jpg1)
  35 assert image1.shape == shape
  36
  37 ys, xs, cs = shape
  38 xb = (xs // BLOCK_SIZE - BLOCKS) // 2
  39 yb = (ys // BLOCK_SIZE - BLOCKS) // 2
  40 print('xb', xb, 'yb', yb)
  41
  42 def bandpass(image):
  43   _, _, cs = image.shape
  44   return numpy.stack(
  45     [
  46       scipy.ndimage.gaussian_filter(image[:, :, i], CUTOFF1, mode = 'mirror') -
  47         scipy.ndimage.gaussian_filter(image[:, :, i], CUTOFF0, mode = 'mirror')
  48       for i in range(cs)
  49     ],
  50     2
  51   )
  52
  53 def correlate(image0_bp, image1_bp, xc, yc):
  54   x0 = xc - BLOCKS * BLOCK_SIZE // 2
  55   y0 = yc - BLOCKS * BLOCK_SIZE // 2
  56   x1 = xc - (BLOCKS + 1) * BLOCK_SIZE // 2
  57   y1 = yc - (BLOCKS + 1) * BLOCK_SIZE // 2
  58
  59   block0 = image0_bp[
  60     y0:y0 + BLOCK_SIZE * BLOCKS,
  61     x0:x0 + BLOCK_SIZE * BLOCKS,
  62     :
  63   ]
  64   block1 = image1_bp[
  65     y1:y1 + (BLOCKS + 1) * BLOCK_SIZE,
  66     x1:x1 + (BLOCKS + 1) * BLOCK_SIZE,
  67     :
  68   ]
  69
  70   # note: swapping block1, block0 flips the output (subtracts x and y
  71   # from BLOCK_SIZE) and we need this to find matching part of image1
  72   corr = numpy.sum(
  73     numpy.stack(
  74       [
  75         scipy.signal.correlate(
  76           block1[:, :, i],
  77           block0[:, :, i],
  78           mode = 'valid'
  79         )
  80         for i in range(cs)
  81       ],
  82       0
  83     ),
  84     0
  85   )
  86   #temp = corr - numpy.mean(corr)
  87   #temp /= 10. * numpy.sqrt(numpy.mean(numpy.square(temp)))
  88   #gamma.write_image(f'corr_{i:d}_{j:d}.jpg', temp + .5)
  89
  90   y, x = numpy.unravel_index(numpy.argmax(corr), corr.shape)
  91   return (
  92     (x - BLOCK_SIZE // 2, y - BLOCK_SIZE // 2)
  93   if (
  94     x >= CUTOFF1 and
  95       x <= BLOCK_SIZE - CUTOFF1 and
  96       y >= CUTOFF1 and
  97       y <= BLOCK_SIZE - CUTOFF1
  98   ) else
  99     None
 100   )
 101
 102 print('bp0')
 103 image0_bp = bandpass(image0)
 104 gamma.write_image('image0_bp.jpg', image0_bp + .5)
 105
 106 print('bp1')
 107 image1_bp = bandpass(image1)
 108 gamma.write_image('image1_bp.jpg', image1_bp + .5)
 109
 110 print('find corner candidates')
 111 buckets = [[[], []], [[], []]]
 112 for i in range(yb - BLOCKS):
 113   for j in range(xb - BLOCKS):
 114     xc = j * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
 115
 116 p_all = []
 117 q_all = []
 118 corner_candidates = []
 119 for i in range(2):
 120   for j in range(2):
 121     # correlate blocks in (i, j)-corner
 122     offsets = []
 123     blocks = []
 124     for k in range(yb):
 125       yc = k * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
 126       if i:
 127         yc = ys - yc
 128       for l in range(xb):
 129         xc = l * BLOCK_SIZE + (BLOCKS + 1) * BLOCK_SIZE // 2
 130         if j:
 131           xc = xs - xc
 132         offset = correlate(image0_bp, image1_bp, xc, yc)
 133         if offset is not None:
 134           offsets.append(offset)
 135           x, y = offset
 136           xc1 = xc + x
 137           yc1 = yc + y
 138           #print('i', i, 'j', j, 'k', k, 'l', l, 'x', x, 'y', y)
 139           p_all.append(numpy.array([xc, yc], numpy.double))
 140           q_all.append(numpy.array([xc1, yc1], numpy.double))
 141           blocks.append((xc, yc, xc1, yc1))
 142
 143     # find the trend in (i, j)-corner
 144     k = len(blocks) // 2
 145     xm = sorted([x for x, _ in offsets])[k]
 146     ym = sorted([y for _, y in offsets])[k]
 147     #print('i', i, 'j', j, 'xm', xm, 'ym', ym)
 148
 149     # choose CORNER_CANDIDATES blocks closest to trend in (i, j)-corner
 150     u = numpy.array(offsets, numpy.double)
 151     v = numpy.array([xm, ym], numpy.double)
 152     dist = numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1)
 153     corner_candidates.append(
 154       sorted(
 155         [(dist[i],) + blocks[i] for i in range(len(blocks))]
 156       )[:CORNER_CANDIDATES]
 157     )
 158 p_all = numpy.stack(p_all, 1)
 159 q_all = numpy.stack(q_all, 1)
 160
 161 # try all combinations of the corner candidates
 162 print('try corner candidates')
 163 p = numpy.zeros((2, 4), numpy.double)
 164 q = numpy.zeros((2, 4), numpy.double)
 165 best_dist = None
 166 best_A = None
 167 for _, xc00, yc00, xc10, yc10 in corner_candidates[0]:
 168   p[0, 0] = xc00
 169   p[1, 0] = yc00
 170   q[0, 0] = xc10
 171   q[1, 0] = yc10
 172   for _, xc01, yc01, xc11, yc11 in corner_candidates[1]:
 173     p[0, 1] = xc01
 174     p[1, 1] = yc01
 175     q[0, 1] = xc11
 176     q[1, 1] = yc11
 177     for _, xc02, yc02, xc12, yc12 in corner_candidates[2]:
 178       p[0, 2] = xc02
 179       p[1, 2] = yc02
 180       q[0, 2] = xc12
 181       q[1, 2] = yc12
 182       for _, xc03, yc03, xc13, yc13 in corner_candidates[3]:
 183         p[0, 3] = xc03
 184         p[1, 3] = yc03
 185         q[0, 3] = xc13
 186         q[1, 3] = yc13
 187
 188         A = perspective.calc_transform(p, q)
 189         dist = numpy.sum(
 190           numpy.square(
 191             q_all - perspective.apply_transform_multi(A, p_all)
 192           )
 193         )
 194         if best_dist is None or dist < best_dist:
 195           best_dist = dist
 196           best_A = A
 197
 198 print('remap')
 199 out_image1 = perspective.remap_image(best_A, image1)
 200
 201 sys.stderr.write(f'write {out_jpg1:s}\n')
 202 gamma.write_image(out_jpg1, out_image1)