Use estimated position of feature within block rather than just block's centre
authorNick Downing <nick@ndcode.org>
Wed, 9 Feb 2022 04:59:48 +0000 (15:59 +1100)
committerNick Downing <nick@ndcode.org>
Wed, 9 Feb 2022 04:59:48 +0000 (15:59 +1100)
correlate.py

index 0df159c..c3be2cb 100755 (executable)
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import gamma
+import math
 import numpy
 import perspective
 import scipy
@@ -26,28 +27,9 @@ CORNER_CANDIDATES = 8
 CUTOFF0 = 64
 CUTOFF1 = 4
 
-in_jpg0 = 'tank_battle/down_2364.jpg'
-in_jpg1 = 'tank_battle/down_2365.jpg'
-out_jpg0 = 'out0_2364.jpg'
-out_jpg1 = 'out0_2365.jpg'
-
-print(f'read {in_jpg0:s}')
-image0 = gamma.read_image(in_jpg0)
-shape = image0.shape
-
-sys.stderr.write(f'write {out_jpg0:s}\n')
-gamma.write_image(out_jpg0, image0)
+EPSILON = 1e-6
 
-print(f'read {in_jpg1:s}')
-image1 = gamma.read_image(in_jpg1)
-assert image1.shape == shape
-
-ys, xs, cs = shape
-xb = (xs // 2 - XM - 2 * XS) // XP
-yb = (ys // 2 - YM - 2 * YS) // YP
-print('xb', xb, 'yb', yb)
-
-def bandpass(image):
+def calc_bandpass(image):
   _, _, cs = image.shape
   return numpy.stack(
     [
@@ -58,18 +40,18 @@ def bandpass(image):
     2
   )
 
-def correlate(image0_bp, image1_bp, xc, yc):
+def calc_match(bandpass0, bandpass1, xc, yc):
   x0 = xc - XM // 2
   y0 = yc - YM // 2
   x1 = xc - XM // 2 - XS
   y1 = yc - YM // 2 - YS
 
-  block0 = image0_bp[
+  block0 = bandpass0[
     y0:y0 + YM,
     x0:x0 + XM,
     :
   ]
-  block1 = image1_bp[
+  block1 = bandpass1[
     y1:y1 + YM + YS * 2,
     x1:x1 + XM + XS * 2,
     :
@@ -95,25 +77,96 @@ def correlate(image0_bp, image1_bp, xc, yc):
   #temp /= 10. * numpy.sqrt(numpy.mean(numpy.square(temp)))
   #gamma.write_image(f'corr_{i:d}_{j:d}.jpg', temp + .5)
 
-  y, x = numpy.unravel_index(numpy.argmax(corr), corr.shape)
-  return (
-    (x - XS, y - YS)
+  # find slippage from correlation
+  yo, xo = numpy.unravel_index(numpy.argmax(corr), corr.shape)
+  max_corr = corr[yo, xo]
   if (
-    x >= CUTOFF1 and
-      x <= XS * 2 - CUTOFF1 and
-      y >= CUTOFF1 and
-      y <= YS * 2 - CUTOFF1
-  ) else
-    None
-  )
+    max_corr < EPSILON or
+      xo < CUTOFF1 or
+      xo > XS * 2 - CUTOFF1 or
+      yo < CUTOFF1 or
+      yo > YS * 2 - CUTOFF1
+  ):
+    return None
 
-print('bp0')
-image0_bp = bandpass(image0)
-gamma.write_image('image0_bp.jpg', image0_bp + .5)
+  # estimate position within block of feature being matched
+  block1 = block1[yo:yo + YM, xo:xo + XM, :]
+  match = numpy.sum(block0 * block1, 2)
+  #print('xxx', numpy.sum(match), max_corr)
+  #assert False
+  xf = (
+    numpy.sum(match, 0) @ numpy.arange(XM, dtype = numpy.double)
+  ) / max_corr
+  yf = (
+    numpy.sum(match, 1) @ numpy.arange(YM, dtype = numpy.double)
+  ) / max_corr
+  #if diag:
+  #  x = int(math.floor(xf))
+  #  y = int(math.floor(yf))
 
-print('bp1')
-image1_bp = bandpass(image1)
-gamma.write_image('image1_bp.jpg', image1_bp + .5)
+  #  diag0 = block0 + .5
+  #  diag0[
+  #    max(y - 21, 0):max(y + 22, 0),
+  #    max(x - 1, 0):max(x + 2, 0),
+  #    :
+  #  ] = 0.
+  #  diag0[
+  #    max(y - 1, 0):max(y + 2, 0),
+  #    max(x - 21, 0):max(x + 22, 0),
+  #    :
+  #  ] = 0.
+  #  gamma.write_image(f'diag_{xc:d}_{yc:d}_0.jpg', diag0)
+
+  #  diag1 = block1 + .5 
+  #  diag1[
+  #    max(y - 21, 0):max(y + 22, 0),
+  #    max(x - 1, 0):max(x + 2, 0),
+  #    :
+  #  ] = 0.
+  #  diag1[
+  #    max(y - 1, 0):max(y + 2, 0),
+  #    max(x - 21, 0):max(x + 22, 0),
+  #    :
+  #  ] = 0.
+  #  gamma.write_image(f'diag_{xc:d}_{yc:d}_1.jpg', diag1)
+
+  # return offset and feature relative to block centre
+  return xo - XS, yo - YS, xf - XM // 2, yf - YM // 2
+
+
+
+diag = False
+if len(sys.argv) >= 2 and sys.argv[1] == '--diag':
+  diag = True
+  del sys.argv[1]
+
+in_jpg0 = 'tank_battle/down_2364.jpg'
+in_jpg1 = 'tank_battle/down_2365.jpg'
+out_jpg0 = 'out0_2364.jpg'
+out_jpg1 = 'out0_2365.jpg'
+
+print(f'read {in_jpg0:s}')
+image0 = gamma.read_image(in_jpg0)
+shape = image0.shape
+
+print('bandpass')
+bandpass0 = calc_bandpass(image0)
+if diag:
+  gamma.write_image('bandpass0.jpg', bandpass0 + .5)
+
+print(f'read {in_jpg1:s}')
+image1 = gamma.read_image(in_jpg1)
+assert image1.shape == shape
+
+print('bandpass')
+bandpass1 = calc_bandpass(image1)
+if diag:
+  gamma.write_image('bandpass1.jpg', bandpass1 + .5)
+
+ys, xs, cs = shape
+xb = (xs // 2 - XM - 2 * XS) // XP
+yb = (ys // 2 - YM - 2 * YS) // YP
+print('xb', xb, 'yb', yb)
 
 print('find corner candidates')
 p_all = []
@@ -121,6 +174,8 @@ q_all = []
 corner_candidates = []
 for i in range(2):
   for j in range(2):
+    print('i', i, 'j', j)
+
     # correlate blocks in (i, j)-corner
     offsets = []
     blocks = []
@@ -132,26 +187,27 @@ for i in range(2):
         xc = XS + XM // 2 + l * XP
         if j:
           xc = xs - xc
-        offset = correlate(image0_bp, image1_bp, xc, yc)
-        if offset is not None:
-          offsets.append(offset)
-          x, y = offset
-          xc1 = xc + x
-          yc1 = yc + y
-          #print('i', i, 'j', j, 'k', k, 'l', l, 'x', x, 'y', y)
-          p_all.append(numpy.array([xc, yc], numpy.double))
-          q_all.append(numpy.array([xc1, yc1], numpy.double))
-          blocks.append((xc, yc, xc1, yc1))
-
-    # find the trend in (i, j)-corner
+        match = calc_match(bandpass0, bandpass1, xc, yc)
+        if match is not None:
+          offsets.append(match [:2])
+          xo, yo, xf, yf = match
+          xf0 = xc + xf
+          yf0 = yc + yf
+          xf1 = xf0 + xo
+          yf1 = yf0 + yo
+          p_all.append(numpy.array([xf0, yf0], numpy.double))
+          q_all.append(numpy.array([xf1, yf1], numpy.double))
+          blocks.append((xf0, yf0, xf1, yf1))
+
+    # find the offset trend (median offset per axis) in (i, j)-corner
     k = len(blocks) // 2
-    xm = sorted([x for x, _ in offsets])[k]
-    ym = sorted([y for _, y in offsets])[k]
-    #print('i', i, 'j', j, 'xm', xm, 'ym', ym)
+    xo_median = sorted([xo for xo, _ in offsets])[k]
+    yo_median = sorted([yo for _, yo in offsets])[k]
+    #print('i', i, 'j', j, 'xo_median', xo_median, 'yo_median', yo_median)
 
     # choose CORNER_CANDIDATES blocks closest to trend in (i, j)-corner
     u = numpy.array(offsets, numpy.double)
-    v = numpy.array([xm, ym], numpy.double)
+    v = numpy.array([xo_median, yo_median], numpy.double)
     dist = numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1)
     corner_candidates.append(
       sorted(
@@ -167,26 +223,27 @@ p = numpy.zeros((2, 4), numpy.double)
 q = numpy.zeros((2, 4), numpy.double)
 best_dist = None
 best_A = None
-for _, xc00, yc00, xc10, yc10 in corner_candidates[0]:
-  p[0, 0] = xc00
-  p[1, 0] = yc00
-  q[0, 0] = xc10
-  q[1, 0] = yc10
-  for _, xc01, yc01, xc11, yc11 in corner_candidates[1]:
-    p[0, 1] = xc01
-    p[1, 1] = yc01
-    q[0, 1] = xc11
-    q[1, 1] = yc11
-    for _, xc02, yc02, xc12, yc12 in corner_candidates[2]:
-      p[0, 2] = xc02
-      p[1, 2] = yc02
-      q[0, 2] = xc12
-      q[1, 2] = yc12
-      for _, xc03, yc03, xc13, yc13 in corner_candidates[3]:
-        p[0, 3] = xc03
-        p[1, 3] = yc03
-        q[0, 3] = xc13
-        q[1, 3] = yc13
+best_p = None # for diag
+for _, xf00, yf00, xf10, yf10 in corner_candidates[0]:
+  p[0, 0] = xf00
+  p[1, 0] = yf00
+  q[0, 0] = xf10
+  q[1, 0] = yf10
+  for _, xf01, yf01, xf11, yf11 in corner_candidates[1]:
+    p[0, 1] = xf01
+    p[1, 1] = yf01
+    q[0, 1] = xf11
+    q[1, 1] = yf11
+    for _, xf02, yf02, xf12, yf12 in corner_candidates[2]:
+      p[0, 2] = xf02
+      p[1, 2] = yf02
+      q[0, 2] = xf12
+      q[1, 2] = yf12
+      for _, xf03, yf03, xf13, yf13 in corner_candidates[3]:
+        p[0, 3] = xf03
+        p[1, 3] = yf03
+        q[0, 3] = xf13
+        q[1, 3] = yf13
 
         A = perspective.calc_transform(p, q)
         dist = numpy.sum(
@@ -197,9 +254,38 @@ for _, xc00, yc00, xc10, yc10 in corner_candidates[0]:
         if best_dist is None or dist < best_dist:
           best_dist = dist
           best_A = A
+          best_p = numpy.copy(p) # for diag
 
 print('remap')
 out_image1 = perspective.remap_image(best_A, image1)
+if diag:
+  for i in range(4):
+    xf0, yf0 = numpy.floor(best_p[:, i]).astype(numpy.int32)
+
+    image0[
+      max(yf0 - 21, 0):max(yf0 + 22, 0),
+      max(xf0 - 1, 0):max(xf0 + 2, 0),
+      :
+    ] = 0.
+    image0[
+      max(yf0 - 1, 0):max(yf0 + 2, 0),
+      max(xf0 - 21, 0):max(xf0 + 22, 0),
+      :
+    ] = 0.
+
+    out_image1[
+      max(yf0 - 21, 0):max(yf0 + 22, 0),
+      max(xf0 - 1, 0):max(xf0 + 2, 0),
+      :
+    ] = 0.
+    out_image1[
+      max(yf0 - 1, 0):max(yf0 + 2, 0),
+      max(xf0 - 21, 0):max(xf0 + 22, 0),
+      :
+    ] = 0.
+
+sys.stderr.write(f'write {out_jpg0:s}\n')
+gamma.write_image(out_jpg0, image0)
 
 sys.stderr.write(f'write {out_jpg1:s}\n')
 gamma.write_image(out_jpg1, out_image1)