From 54f4e699558e332ccb6ea1a61680c2458f092bd9 Mon Sep 17 00:00:00 2001 From: Nick Downing Date: Wed, 9 Feb 2022 17:19:24 +1100 Subject: [PATCH] Implement batch processing, tune parameters slightly --- correlate.py | 306 +++++++++++++++++++++++++++------------------------ 1 file changed, 163 insertions(+), 143 deletions(-) diff --git a/correlate.py b/correlate.py index c3be2cb..43bbb6b 100755 --- a/correlate.py +++ b/correlate.py @@ -3,6 +3,7 @@ import gamma import math import numpy +import os import perspective import scipy import scipy.ndimage @@ -22,7 +23,7 @@ YP = 64 XS = 64 YS = 64 -CORNER_CANDIDATES = 8 +CORNER_CANDIDATES = 16 CUTOFF0 = 64 CUTOFF1 = 4 @@ -117,7 +118,7 @@ def calc_match(bandpass0, bandpass1, xc, yc): # ] = 0. # gamma.write_image(f'diag_{xc:d}_{yc:d}_0.jpg', diag0) - # diag1 = block1 + .5 + # diag1 = block1 + .5 # diag1[ # max(y - 21, 0):max(y + 22, 0), # max(x - 1, 0):max(x + 2, 0), @@ -133,159 +134,178 @@ def calc_match(bandpass0, bandpass1, xc, yc): # 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] +if len(sys.argv) < 3: + print(f'usage: {sys.argv[0]} in_prefix out_prefix') + sys.exit(1) +in_prefix = sys.argv[1] +out_prefix = sys.argv[2] + +# compile list of input and output files by directory search +i = in_prefix.rfind('/') + 1 +in_prefix0 = in_prefix[:i] # directory part +in_prefix1 = in_prefix[i:] # filename part +files = [ + (file, out_prefix + file[len(in_prefix):]) + for file in sorted( + [ + in_prefix0 + file + for file in os.listdir(in_prefix0 if len(in_prefix0) else '.') + if file[:len(in_prefix1)] == in_prefix1 + ] + ) +] -in_jpg0 = 'tank_battle/down_2364.jpg' -in_jpg1 = 'tank_battle/down_2365.jpg' -out_jpg0 = 'out0_2364.jpg' -out_jpg1 = 'out0_2365.jpg' +# first file is special as no transformation needs to be done +in_jpg, out_jpg = files[0] -print(f'read {in_jpg0:s}') -image0 = gamma.read_image(in_jpg0) +print(f'read {in_jpg:s}') +image0 = gamma.read_image(in_jpg) 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) +sys.stderr.write(f'write {out_jpg:s}\n') +gamma.write_image(out_jpg, image0) 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 = [] -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 = [] - for k in range(yb): - yc = YS + YM // 2 + k * YP - if i: - yc = ys - yc - for l in range(xb): - xc = XS + XM // 2 + l * XP - if j: - xc = xs - xc - 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 - 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([xo_median, yo_median], numpy.double) - dist = numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1) - corner_candidates.append( - sorted( - [(dist[i],) + blocks[i] for i in range(len(blocks))] - )[:CORNER_CANDIDATES] - ) -p_all = numpy.stack(p_all, 1) -q_all = numpy.stack(q_all, 1) - -# try all combinations of the corner candidates -print('try corner candidates') -p = numpy.zeros((2, 4), numpy.double) -q = numpy.zeros((2, 4), numpy.double) -best_dist = None -best_A = None -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( - numpy.square( - q_all - perspective.apply_transform_multi(A, p_all) - ) - ) - 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) +print('bandpass') +bandpass0 = calc_bandpass(image0) 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) + gamma.write_image('bandpass0.jpg', bandpass0 + .5) + +# loop through remaining files comparing each to the previous +cumulative_A = numpy.identity(3, numpy.double) +for file in range(1, len(files)): + in_jpg, out_jpg = files[file] + + print(f'read {in_jpg:s}') + image1 = gamma.read_image(in_jpg) + assert image1.shape == shape + + print('bandpass') + bandpass1 = calc_bandpass(image1) + if diag: + gamma.write_image(f'bandpass{file:d}.jpg', bandpass1 + .5) + + print('find corner candidates') + p_all = [] + 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 = [] + for k in range(yb): + yc = YS + YM // 2 + k * YP + if i: + yc = ys - yc + for l in range(xb): + xc = XS + XM // 2 + l * XP + if j: + xc = xs - xc + 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 + 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([xo_median, yo_median], numpy.double) + dist = numpy.sum(numpy.square(u - v[numpy.newaxis, :]), 1) + corner_candidates.append( + sorted( + [(dist[i],) + blocks[i] for i in range(len(blocks))] + )[:CORNER_CANDIDATES] + ) + p_all = numpy.stack(p_all, 1) + q_all = numpy.stack(q_all, 1) + + # try all combinations of the corner candidates + print('try corner candidates') + p = numpy.zeros((2, 4), numpy.double) + q = numpy.zeros((2, 4), numpy.double) + best_dist = None + best_A = None + 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 = list( + numpy.sum( + numpy.square( + q_all - perspective.apply_transform_multi(A, p_all) + ), + 0 + ) + ) + dist = sorted(dist)[len(dist) // 2] # median position error + + if best_dist is None or dist < best_dist: + best_dist = dist + best_A = A + best_p = numpy.copy(p) # for diag + cumulative_A = best_A @ cumulative_A + + print('remap') + out_image1 = perspective.remap_image(cumulative_A, image1) + if diag: + for i in range(4): + xf0, yf0 = numpy.floor(best_p[:, i]).astype(numpy.int32) + + 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_jpg:s}\n') + gamma.write_image(out_jpg, out_image1) + + image0 = image1 + bandpass0 = bandpass1 -- 2.34.1