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