correlate.py

   1 #!/usr/bin/env python3
   2
   3 import gamma
   4 import math
   5 import numpy
   6 import os
   7 import perspective
   8 import scipy
   9 import scipy.ndimage
  10 import scipy.signal
  11 import sys
  12
  13 # size of block that will be matched
  14 # (correlate a block this size against a block with added slippage all around)
  15 XM = 128
  16 YM = 128
  17
  18 # pitch between the block centres
  19 XP = 64
  20 YP = 64
  21
  22 # allowable +/- slippage between pairs
  23 XS = 64
  24 YS = 64
  25
  26 CORNER_CANDIDATES = 16
  27
  28 CUTOFF0 = 64
  29 CUTOFF1 = 4
  30
  31 EPSILON = 1e-6
  32
  33 def calc_bandpass(image):
  34   _, _, cs = image.shape
  35   return numpy.stack(
  36     [
  37       scipy.ndimage.gaussian_filter(image[:, :, i], CUTOFF1, mode = 'mirror') -
  38         scipy.ndimage.gaussian_filter(image[:, :, i], CUTOFF0, mode = 'mirror')
  39       for i in range(cs)
  40     ],
  41     2
  42   )
  43
  44 def calc_match(bandpass0, bandpass1, xc, yc):
  45   x0 = xc - XM // 2
  46   y0 = yc - YM // 2
  47   x1 = xc - XM // 2 - XS
  48   y1 = yc - YM // 2 - YS
  49
  50   block0 = bandpass0[
  51     y0:y0 + YM,
  52     x0:x0 + XM,
  53     :
  54   ]
  55   block1 = bandpass1[
  56     y1:y1 + YM + YS * 2,
  57     x1:x1 + XM + XS * 2,
  58     :
  59   ]
  60
  61   # note: swapping block1, block0 flips the output (subtracts x and y
  62   # from BLOCK_SIZE) and we need this to find matching part of image1
  63   corr = numpy.sum(
  64     numpy.stack(
  65       [
  66         scipy.signal.correlate(
  67           block1[:, :, i],
  68           block0[:, :, i],
  69           mode = 'valid'
  70         )
  71         for i in range(cs)
  72       ],
  73       0
  74     ),
  75     0
  76   )
  77   #temp = corr - numpy.mean(corr)
  78   #temp /= 10. * numpy.sqrt(numpy.mean(numpy.square(temp)))
  79   #gamma.write_image(f'corr_{i:d}_{j:d}.jpg', temp + .5)
  80
  81   # find slippage from correlation
  82   yo, xo = numpy.unravel_index(numpy.argmax(corr), corr.shape)
  83   max_corr = corr[yo, xo]
  84   if (
  85     max_corr < EPSILON or
  86       xo < CUTOFF1 or
  87       xo > XS * 2 - CUTOFF1 or
  88       yo < CUTOFF1 or
  89       yo > YS * 2 - CUTOFF1
  90   ):
  91     return None
  92
  93   # estimate position within block of feature being matched
  94   block1 = block1[yo:yo + YM, xo:xo + XM, :]
  95   match = numpy.sum(block0 * block1, 2)
  96   #print('xxx', numpy.sum(match), max_corr)
  97   #assert False
  98   xf = (
  99     numpy.sum(match, 0) @ numpy.arange(XM, dtype = numpy.double)
 100   ) / max_corr
 101   yf = (
 102     numpy.sum(match, 1) @ numpy.arange(YM, dtype = numpy.double)
 103   ) / max_corr
 104   #if diag:
 105   #  x = int(math.floor(xf))
 106   #  y = int(math.floor(yf))
 107
 108   #  diag0 = block0 + .5
 109   #  diag0[
 110   #    max(y - 21, 0):max(y + 22, 0),
 111   #    max(x - 1, 0):max(x + 2, 0),
 112   #    :
 113   #  ] = 0.
 114   #  diag0[
 115   #    max(y - 1, 0):max(y + 2, 0),
 116   #    max(x - 21, 0):max(x + 22, 0),
 117   #    :
 118   #  ] = 0.
 119   #  gamma.write_image(f'diag_{xc:d}_{yc:d}_0.jpg', diag0)
 120
 121   #  diag1 = block1 + .5
 122   #  diag1[
 123   #    max(y - 21, 0):max(y + 22, 0),
 124   #    max(x - 1, 0):max(x + 2, 0),
 125   #    :
 126   #  ] = 0.
 127   #  diag1[
 128   #    max(y - 1, 0):max(y + 2, 0),
 129   #    max(x - 21, 0):max(x + 22, 0),
 130   #    :
 131   #  ] = 0.
 132   #  gamma.write_image(f'diag_{xc:d}_{yc:d}_1.jpg', diag1)
 133
 134   # return offset and feature relative to block centre
 135   return xo - XS, yo - YS, xf - XM // 2, yf - YM // 2
 136
 137 diag = False
 138 if len(sys.argv) >= 2 and sys.argv[1] == '--diag':
 139   diag = True
 140   del sys.argv[1]
 141 if len(sys.argv) < 3:
 142   print(f'usage: {sys.argv[0]} in_prefix out_prefix')
 143   sys.exit(1)
 144 in_prefix = sys.argv[1]
 145 out_prefix = sys.argv[2]
 146
 147 # compile list of input and output files by directory search
 148 i = in_prefix.rfind('/') + 1
 149 in_prefix0 = in_prefix[:i] # directory part
 150 in_prefix1 = in_prefix[i:] # filename part
 151 files = [
 152   (file, out_prefix + file[len(in_prefix):])
 153   for file in sorted(
 154     [
 155       in_prefix0 + file
 156       for file in os.listdir(in_prefix0 if len(in_prefix0) else '.')
 157       if file[:len(in_prefix1)] == in_prefix1
 158     ]
 159   )
 160 ]
 161
 162 # first file is special as no transformation needs to be done
 163 in_jpg, out_jpg = files[0]
 164
 165 print(f'read {in_jpg:s}')
 166 image0 = gamma.read_image(in_jpg)
 167 shape = image0.shape
 168
 169 sys.stderr.write(f'write {out_jpg:s}\n')
 170 gamma.write_image(out_jpg, image0)
 171
 172 ys, xs, cs = shape
 173 xb = (xs // 2 - XM - 2 * XS) // XP
 174 yb = (ys // 2 - YM - 2 * YS) // YP
 175 print('xb', xb, 'yb', yb)
 176
 177 print('bandpass')
 178 bandpass0 = calc_bandpass(image0)
 179 if diag:
 180   gamma.write_image('bandpass0.jpg', bandpass0 + .5)
 181
 182 # loop through remaining files comparing each to the previous
 183 cumulative_A = numpy.identity(3, numpy.double)
 184 for file in range(1, len(files)):
 185   in_jpg, out_jpg = files[file]
 186
 187   print(f'read {in_jpg:s}')
 188   image1 = gamma.read_image(in_jpg)
 189   assert image1.shape == shape
 190
 191   print('bandpass')
 192   bandpass1 = calc_bandpass(image1)
 193   if diag:
 194     gamma.write_image(f'bandpass{file:d}.jpg', bandpass1 + .5)
 195
 196   print('find corner candidates')
 197   p_all = []
 198   q_all = []
 199   corner_candidates = []
 200   for i in range(2):
 201     for j in range(2):
 202       print('i', i, 'j', j)
 203
 204       # correlate blocks in (i, j)-corner
 205       offsets = []
 206       blocks = []
 207       for k in range(yb):
 208         yc = YS + YM // 2 + k * YP
 209         if i:
 210           yc = ys - yc
 211         for l in range(xb):
 212           xc = XS + XM // 2 + l * XP
 213           if j:
 214             xc = xs - xc
 215           match = calc_match(bandpass0, bandpass1, xc, yc)
 216           if match is not None:
 217             offsets.append(match [:2])
 218             xo, yo, xf, yf = match
 219             xf0 = xc + xf
 220             yf0 = yc + yf
 221             xf1 = xf0 + xo
 222             yf1 = yf0 + yo
 223             p_all.append(numpy.array([xf0, yf0], numpy.double))
 224             q_all.append(numpy.array([xf1, yf1], numpy.double))
 225             blocks.append((xf0, yf0, xf1, yf1))
 226
 227       # find the offset trend (median offset per axis) in (i, j)-corner
 228       k = len(blocks) // 2
 229       xo_median = sorted([xo for xo, _ in offsets])[k]
 230       yo_median = sorted([yo for _, yo in offsets])[k]
 231       #print('i', i, 'j', j, 'xo_median', xo_median, 'yo_median', yo_median)
 232
 233       # choose CORNER_CANDIDATES blocks closest to trend in (i, j)-corner
 234       u = numpy.array(offsets, numpy.double)
 235       v = numpy.array([xo_median, yo_median], numpy.double)
 236       dist = numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1)
 237       corner_candidates.append(
 238         sorted(
 239           [(dist[i],) + blocks[i] for i in range(len(blocks))]
 240         )[:CORNER_CANDIDATES]
 241       )
 242   p_all = numpy.stack(p_all, 1)
 243   q_all = numpy.stack(q_all, 1)
 244
 245   # try all combinations of the corner candidates
 246   print('try corner candidates')
 247   p = numpy.zeros((2, 4), numpy.double)
 248   q = numpy.zeros((2, 4), numpy.double)
 249   best_dist = None
 250   best_A = None
 251   best_p = None # for diag
 252   for _, xf00, yf00, xf10, yf10 in corner_candidates[0]:
 253     p[0, 0] = xf00
 254     p[1, 0] = yf00
 255     q[0, 0] = xf10
 256     q[1, 0] = yf10
 257     for _, xf01, yf01, xf11, yf11 in corner_candidates[1]:
 258       p[0, 1] = xf01
 259       p[1, 1] = yf01
 260       q[0, 1] = xf11
 261       q[1, 1] = yf11
 262       for _, xf02, yf02, xf12, yf12 in corner_candidates[2]:
 263         p[0, 2] = xf02
 264         p[1, 2] = yf02
 265         q[0, 2] = xf12
 266         q[1, 2] = yf12
 267         for _, xf03, yf03, xf13, yf13 in corner_candidates[3]:
 268           p[0, 3] = xf03
 269           p[1, 3] = yf03
 270           q[0, 3] = xf13
 271           q[1, 3] = yf13
 272
 273           A = perspective.calc_transform(p, q)
 274           dist = list(
 275             numpy.sum(
 276               numpy.square(
 277                 q_all - perspective.apply_transform_multi(A, p_all)
 278               ),
 279               0
 280             )
 281           )
 282           dist = sorted(dist)[len(dist) // 2] # median position error
 283
 284           if best_dist is None or dist < best_dist:
 285             best_dist = dist
 286             best_A = A
 287             best_p = numpy.copy(p) # for diag
 288   cumulative_A = best_A @ cumulative_A
 289
 290   print('remap')
 291   out_image1 = perspective.remap_image(cumulative_A, image1)
 292   if diag:
 293     for i in range(4):
 294       xf0, yf0 = numpy.floor(best_p[:, i]).astype(numpy.int32)
 295
 296       out_image1[
 297         max(yf0 - 21, 0):max(yf0 + 22, 0),
 298         max(xf0 - 1, 0):max(xf0 + 2, 0),
 299         :
 300       ] = 0.
 301       out_image1[
 302         max(yf0 - 1, 0):max(yf0 + 2, 0),
 303         max(xf0 - 21, 0):max(xf0 + 22, 0),
 304         :
 305       ] = 0.
 306
 307   sys.stderr.write(f'write {out_jpg:s}\n')
 308   gamma.write_image(out_jpg, out_image1)
 309
 310   image0 = image1
 311   bandpass0 = bandpass1