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 print('bp0')
  41 image0_lp0 = numpy.stack(
  42   [
  43     scipy.ndimage.gaussian_filter(image0[:, :, i], CUTOFF0, mode = 'mirror')
  44     for i in range(cs)
  45   ],
  46   2
  47 )
  48 image0_lp1 = numpy.stack(
  49   [
  50     scipy.ndimage.gaussian_filter(image0[:, :, i], CUTOFF1, mode = 'mirror')
  51     for i in range(cs)
  52   ],
  53   2
  54 )
  55 image0_bp = image0_lp1 - image0_lp0
  56 gamma.write_image('image0_bp.jpg', image0_bp + .5)
  57
  58 print('bp1')
  59 image1_lp0 = numpy.stack(
  60   [
  61     scipy.ndimage.gaussian_filter(image1[:, :, i], CUTOFF0, mode = 'mirror')
  62     for i in range(cs)
  63   ],
  64   2
  65 )
  66 image1_lp1 = numpy.stack(
  67   [
  68     scipy.ndimage.gaussian_filter(image1[:, :, i], CUTOFF1, mode = 'mirror')
  69     for i in range(cs)
  70   ],
  71   2
  72 )
  73 image1_bp = image1_lp1 - image1_lp0
  74 gamma.write_image('image1_bp.jpg', image1_bp + .5)
  75
  76 print('align')
  77 buckets = [[[], []], [[], []]]
  78 for i in range(yb - BLOCKS):
  79   y1 = i * BLOCK_SIZE
  80   y0 = y1 + BLOCK_SIZE // 2
  81   for j in range(xb - BLOCKS):
  82     x1 = j * BLOCK_SIZE
  83     x0 = x1 + BLOCK_SIZE // 2
  84
  85     block0 = image0_bp[
  86       y0:y0 + BLOCK_SIZE * BLOCKS,
  87       x0:x0 + BLOCK_SIZE * BLOCKS,
  88       :
  89     ]
  90     block1 = image1_bp[
  91       y1:y1 + BLOCK_SIZE * (BLOCKS + 1),
  92       x1:x1 + BLOCK_SIZE * (BLOCKS + 1),
  93       :
  94     ]
  95
  96     # note: swapping block1, block0 flips the output (subtracts x and y
  97     # from BLOCK_SIZE) and we need this to find matching part of image1
  98     corr = numpy.sum(
  99       numpy.stack(
 100         [
 101           scipy.signal.correlate(
 102             block1[:, :, i],
 103             block0[:, :, i],
 104             mode = 'valid'
 105           )
 106           for i in range(cs)
 107         ],
 108         0
 109       ),
 110       0
 111     )
 112     #temp = corr - numpy.mean(corr)
 113     #temp /= 10. * numpy.sqrt(numpy.mean(numpy.square(temp)))
 114     #gamma.write_image(f'corr_{i:d}_{j:d}.jpg', temp + .5)
 115
 116     y, x = numpy.unravel_index(numpy.argmax(corr), corr.shape)
 117     if (
 118       y >= CUTOFF1 and
 119         y <= BLOCK_SIZE - CUTOFF1 and
 120         x >= CUTOFF1 and
 121         x <= BLOCK_SIZE - CUTOFF1
 122     ):
 123       buckets[i >= (yb - BLOCKS) // 2][j >= (xb - BLOCKS) // 2].append(
 124         (x0, y0, x1, y1, x, y)
 125       )
 126
 127 p = []
 128 q = []
 129 for i in range(2):
 130   for j in range(2):
 131     k = len(buckets[i][j]) // 2
 132     xm = sorted([x for _, _, _, _, x, _ in buckets[i][j]])[k]
 133     ym = sorted([y for _, _, _, _, _, y in buckets[i][j]])[k]
 134     #print('i', i, 'j', j, 'xm', xm, 'ym', ym)
 135     u = numpy.array(
 136       [[x, y] for _, _, _, _, x, y in buckets[i][j]],
 137       numpy.double
 138     )
 139     v = numpy.array([xm, ym], numpy.double)
 140     k = numpy.argmin(numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1))
 141     x0, y0, x1, y1, x, y = buckets[i][j][k]
 142     #print('i', i, 'j', j, 'x', x, 'y', y)
 143     p.append(
 144       numpy.array(
 145         [
 146           x0 + (BLOCKS * BLOCK_SIZE // 2),
 147           y0 + (BLOCKS * BLOCK_SIZE // 2)
 148         ],
 149         numpy.double
 150       )
 151     )
 152     q.append(
 153       numpy.array(
 154         [
 155           x1 + x + (BLOCKS * BLOCK_SIZE // 2),
 156           y1 + y + (BLOCKS * BLOCK_SIZE // 2)
 157         ],
 158         numpy.double
 159       )
 160     )
 161 p = numpy.stack(p, 1)
 162 q = numpy.stack(q, 1)
 163 A = perspective.calc_transform(p, q)
 164
 165 print('remap')
 166 coords = numpy.zeros((2, ys, xs), numpy.double)
 167 coords[0, :, :] = numpy.arange(ys, dtype = numpy.double)[:, numpy.newaxis]
 168 coords[1, :, :] = numpy.arange(xs, dtype = numpy.double)[numpy.newaxis, :]
 169 mapped_coords = perspective.apply_transform_multi(
 170   A,
 171   coords[::-1, :, :].reshape((2, ys * xs))
 172 ).reshape(2, ys, xs)[::-1, :, :]
 173 out_image1 = numpy.stack(
 174   [
 175     scipy.ndimage.map_coordinates(image1[:, :, j], mapped_coords)
 176     for j in range(cs)
 177   ],
 178   2
 179 )
 180
 181 sys.stderr.write(f'write {out_jpg1:s}\n')
 182 gamma.write_image(out_jpg1, out_image1)