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 = 128
  24 YS = 128
  25
  26 CORNER_CANDIDATES = 10
  27
  28 CUTOFF0 = 8
  29 CUTOFF1 = 2
  30
  31 EPSILON = 1e-6
  32
  33 def normalize(image):
  34   _, _, cs = image.shape
  35   num = image - 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   denom = numpy.sqrt(
  44     scipy.ndimage.gaussian_filter(
  45       numpy.sum(numpy.square(num), 2),
  46       CUTOFF0,
  47       mode = 'mirror'
  48     )
  49   )
  50   denom[denom < EPSILON] = EPSILON
  51   return num / denom[:, :, numpy.newaxis]
  52
  53 def calc_match(normalized0, normalized1, xc, yc):
  54   x0 = xc - XM // 2
  55   y0 = yc - YM // 2
  56   x1 = xc - XM // 2 - XS
  57   y1 = yc - YM // 2 - YS
  58
  59   block0 = normalized0[
  60     y0:y0 + YM,
  61     x0:x0 + XM,
  62     :
  63   ]
  64   block1 = normalized1[
  65     y1:y1 + YM + YS * 2,
  66     x1:x1 + XM + XS * 2,
  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   # find slippage from correlation
  87   yo, xo = numpy.unravel_index(numpy.argmax(corr), corr.shape)
  88   max_corr = corr[yo, xo]
  89   if diag:
  90     temp = .5 + .5 * (corr / max_corr)
  91     temp[
  92       max(yo - 21, 0):max(yo + 22, 0),
  93       max(xo - 1, 0):max(xo + 2, 0)
  94     ] = 0.
  95     temp[
  96       max(yo - 1, 0):max(yo + 2, 0),
  97       max(xo - 21, 0):max(xo + 22, 0)
  98     ] = 0.
  99     gamma.write_image(f'diag_{xc:d}_{yc:d}_corr.png', temp)
 100   if (
 101     max_corr < EPSILON or
 102       xo < CUTOFF1 or
 103       xo > XS * 2 - CUTOFF1 or
 104       yo < CUTOFF1 or
 105       yo > YS * 2 - CUTOFF1
 106   ):
 107     return None
 108
 109   # estimate position within block of feature being matched
 110   block10 = block1[yo:yo + YM, xo:xo + XM, :]
 111   match = numpy.sum(block0 * block10, 2)
 112   #print('xxx', numpy.sum(match), max_corr)
 113   #assert False
 114   xf = (
 115     numpy.sum(match, 0) @ numpy.arange(XM, dtype = numpy.double)
 116   ) / max_corr
 117   yf = (
 118     numpy.sum(match, 1) @ numpy.arange(YM, dtype = numpy.double)
 119   ) / max_corr
 120   if diag:
 121     x = int(math.floor(xf))
 122     y = int(math.floor(yf))
 123
 124     temp = match + .5
 125     temp[
 126       max(y - 21, 0):max(y + 22, 0),
 127       max(x - 1, 0):max(x + 2, 0)
 128     ] = 0.
 129     temp[
 130       max(y - 1, 0):max(y + 2, 0),
 131       max(x - 21, 0):max(x + 22, 0)
 132     ] = 0.
 133     gamma.write_image(f'diag_{xc:d}_{yc:d}_match.png', temp)
 134
 135     diag0 = block0 + .5
 136     diag0[
 137       max(y - 21, 0):max(y + 22, 0),
 138       max(x - 1, 0):max(x + 2, 0),
 139       :
 140     ] = 0.
 141     diag0[
 142       max(y - 1, 0):max(y + 2, 0),
 143       max(x - 21, 0):max(x + 22, 0),
 144       :
 145     ] = 0.
 146     gamma.write_image(f'diag_{xc:d}_{yc:d}_block0.png', diag0)
 147
 148     diag1 = block10 + .5
 149     diag1[
 150       max(y - 21, 0):max(y + 22, 0),
 151       max(x - 1, 0):max(x + 2, 0),
 152       :
 153     ] = 0.
 154     diag1[
 155       max(y - 1, 0):max(y + 2, 0),
 156       max(x - 21, 0):max(x + 22, 0),
 157       :
 158     ] = 0.
 159     gamma.write_image(f'diag_{xc:d}_{yc:d}_block10.png', diag1)
 160
 161     x += xo
 162     y += yo
 163     diag1 = block1 + .5
 164     diag1[
 165       max(y - 21, 0):max(y + 22, 0),
 166       max(x - 1, 0):max(x + 2, 0),
 167       :
 168     ] = 0.
 169     diag1[
 170       max(y - 1, 0):max(y + 2, 0),
 171       max(x - 21, 0):max(x + 22, 0),
 172       :
 173     ] = 0.
 174     gamma.write_image(f'diag_{xc:d}_{yc:d}_block11.png', diag1)
 175
 176   # return offset and feature relative to block centre
 177   return xo - XS, yo - YS, xf - XM // 2, yf - YM // 2
 178
 179 diag = False
 180 if len(sys.argv) >= 2 and sys.argv[1] == '--diag':
 181   diag = True
 182   del sys.argv[1]
 183 if len(sys.argv) < 3:
 184   print(f'usage: {sys.argv[0]} in_prefix out_prefix')
 185   sys.exit(1)
 186 in_prefix = sys.argv[1]
 187 out_prefix = sys.argv[2]
 188
 189 # compile list of input and output files by directory search
 190 i = in_prefix.rfind('/') + 1
 191 in_prefix0 = in_prefix[:i] # directory part
 192 in_prefix1 = in_prefix[i:] # filename part
 193 files = [
 194   (file, out_prefix + file[len(in_prefix):])
 195   for file in sorted(
 196     [
 197       in_prefix0 + file
 198       for file in os.listdir(in_prefix0 if len(in_prefix0) else '.')
 199       if file[:len(in_prefix1)] == in_prefix1
 200     ]
 201   )
 202 ]
 203
 204 # first file is special as no transformation needs to be done
 205 in_png, out_png = files[0]
 206
 207 print(f'read {in_png:s}')
 208 image0 = gamma.read_image(in_png)
 209 shape = image0.shape
 210
 211 sys.stderr.write(f'write {out_png:s}\n')
 212 gamma.write_image(out_png, image0)
 213
 214 ys, xs, cs = shape
 215 xb = (xs - XM - 2 * XS) // XP // 2
 216 yb = (ys - YM - 2 * YS) // YP // 2
 217 print('xb', xb, 'yb', yb)
 218
 219 print('normalize')
 220 normalized0 = normalize(image0)
 221 if diag:
 222   gamma.write_image('normalized0.png', normalized0 + .5)
 223
 224 # loop through remaining files comparing each to the previous
 225 cumulative_A = numpy.identity(3, numpy.double)
 226 for file in range(1, len(files)):
 227   in_png, out_png = files[file]
 228
 229   print(f'read {in_png:s}')
 230   image1 = gamma.read_image(in_png)
 231   assert image1.shape == shape
 232
 233   print('normalize')
 234   normalized1 = normalize(image1)
 235   if diag:
 236     gamma.write_image(f'normalized{file:d}.png', normalized1 + .5)
 237
 238   print('find corner candidates')
 239   p_all = []
 240   q_all = []
 241   for i in range(2):
 242     for j in range(2):
 243       print('i', i, 'j', j)
 244
 245       # correlate blocks in (i, j)-corner
 246       offsets = []
 247       blocks = []
 248       for k in range(yb):
 249         yc = YS + YM // 2 + k * YP
 250         if i:
 251           yc = ys - yc
 252         for l in range(xb):
 253           xc = XS + XM // 2 + l * XP
 254           if j:
 255             xc = xs - xc
 256           match = calc_match(normalized0, normalized1, xc, yc)
 257           if match is not None:
 258             offsets.append(match [:2])
 259             xo, yo, xf, yf = match
 260             #xf = 0
 261             #yf = 0
 262             xf0 = xc + xf
 263             yf0 = yc + yf
 264             xf1 = xf0 + xo
 265             yf1 = yf0 + yo
 266             blocks.append((xf0, yf0, xf1, yf1))
 267
 268       # find the offset trend (median offset per axis) in (i, j)-corner
 269       k = len(offsets) // 2
 270       xo_median = sorted([xo for xo, _ in offsets])[k]
 271       yo_median = sorted([yo for _, yo in offsets])[k]
 272       #print('i', i, 'j', j, 'xo_median', xo_median, 'yo_median', yo_median)
 273
 274       # sort all offsets by their closeness to trend in (i,j)-corner
 275       u = numpy.array(offsets, numpy.double)
 276       v = numpy.array([xo_median, yo_median], numpy.double)
 277       dist = numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1)
 278
 279       p = []
 280       q = []
 281       for _, k in sorted([(dist[k], k) for k in range(dist.shape[0])]):
 282         xf0, yf0, xf1, yf1 = blocks[k]
 283         p.append(numpy.array([xf0, yf0], numpy.double))
 284         q.append(numpy.array([xf1, yf1], numpy.double))
 285       p = numpy.stack(p, 1)
 286       q = numpy.stack(q, 1)
 287
 288       p_all.append(p)
 289       q_all.append(q)
 290
 291   # try all combinations of the corner candidates
 292   print('try corner candidates')
 293   best_dist = None
 294   best_A = None
 295   for i in range(min(p_all[0].shape[1], CORNER_CANDIDATES)):
 296     for j in range(min(p_all[1].shape[1], CORNER_CANDIDATES)):
 297       for k in range(min(p_all[2].shape[1], CORNER_CANDIDATES)):
 298         for l in range(min(p_all[3].shape[1], CORNER_CANDIDATES)):
 299           A = perspective.calc_transform(
 300             numpy.stack(
 301               [p_all[0][:, i], p_all[1][:, j], p_all[2][:, k], p_all[3][:, l]],
 302               1
 303             ),
 304             numpy.stack(
 305               [q_all[0][:, i], q_all[1][:, j], q_all[2][:, k], q_all[3][:, l]],
 306               1
 307             )
 308           )
 309           dist = 0.
 310           for m in range(4):
 311             dists = list(
 312               numpy.sum(
 313                 numpy.square(
 314                   q_all[m] - perspective.apply_transform_multi(A, p_all[m])
 315                 ),
 316                 0
 317               )
 318             )
 319             dist += sorted(dists)[len(dists) // 2] # median position error
 320
 321           if best_dist is None or dist < best_dist:
 322             best_dist = dist
 323             best_A = A
 324   cumulative_A = best_A @ cumulative_A
 325
 326   if diag:
 327     diag0 = numpy.copy(image0)
 328     diag1 = numpy.copy(image1)
 329     for i in range(4):
 330       for j in range(p_all[i].shape[1]):
 331         xf0, yf0 = numpy.floor(p_all[i][:, j]).astype(numpy.int32)
 332         diag0[
 333           max(yf0 - 21, 0):max(yf0 + 22, 0),
 334           max(xf0 - 1, 0):max(xf0 + 2, 0),
 335           :
 336         ] = 1.
 337         diag0[
 338           max(yf0 - 1, 0):max(yf0 + 2, 0),
 339           max(xf0 - 21, 0):max(xf0 + 22, 0),
 340           :
 341         ] = 1.
 342
 343         xf1, yf1 = numpy.floor(q_all[i][:, j]).astype(numpy.int32)
 344         diag1[
 345           max(yf1 - 21, 0):max(yf1 + 22, 0),
 346           max(xf1 - 1, 0):max(xf1 + 2, 0),
 347           :
 348         ] = 0.
 349         diag1[
 350           max(yf1 - 1, 0):max(yf1 + 2, 0),
 351           max(xf1 - 21, 0):max(xf1 + 22, 0),
 352           :
 353         ] = 0.
 354     gamma.write_image(f'diag_file{file:d}_0.png', diag0)
 355     gamma.write_image(f'diag_file{file:d}_1.png', diag1)
 356
 357   print('remap')
 358   out_image1 = perspective.remap_image(cumulative_A, image1)
 359
 360   sys.stderr.write(f'write {out_png:s}\n')
 361   gamma.write_image(out_png, out_image1)
 362
 363   image0 = image1
 364   normalized0 = normalized1