align.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 sys
   9
  10 sys.stderr.write(
  11   '''stdin: control blocks description
  12 xs0,ys0 xs1,ys1 xs2,ys2 xs3,ys3 : size (pixels)
  13 xf0,yf0 xf1,yf1 xf2,yf2 xf3,yf3 : freedom (fraction of size)
  14 in.jpg out.jpg x0,y0 x1,y1 x2,y2 x3,y3 : top left (pixels)
  15 ...
  16 '''
  17 )
  18
  19 size = numpy.array(
  20   [
  21     [int(j) for j in i.split(',')]
  22     for i in sys.stdin.readline().split()
  23   ],
  24   numpy.int32
  25 ).transpose()
  26 assert size.shape == (2, 4)
  27 freedom = numpy.array(
  28   [
  29     [float(j) for j in i.split(',')]
  30     for i in sys.stdin.readline().split()
  31   ],
  32   numpy.double
  33 ).transpose()
  34 assert freedom.shape == (2, 4)
  35
  36 files = []
  37 line = sys.stdin.readline()
  38 while len(line):
  39   fields = line.split()
  40   in_jpg = fields[0]
  41   out_jpg = fields[1]
  42   top_left = numpy.array(
  43     [
  44       [int(j) for j in i.split(',')]
  45       for i in fields[2:]
  46     ],
  47     numpy.int32
  48   ).transpose()
  49   assert top_left.shape == (2, 4)
  50   files.append((in_jpg, out_jpg, top_left))
  51   line = sys.stdin.readline()
  52
  53 in_jpg, out_jpg, top_left0 = files[0]
  54 sys.stderr.write(f'read {in_jpg:s}\n')
  55 image0 = gamma.read_image(in_jpg)
  56 shape = image0.shape
  57 sys.stderr.write(f'shape {str(shape):s}\n')
  58
  59 sys.stderr.write(f'write {out_jpg:s}\n')
  60 gamma.write_image(out_jpg, image0)
  61
  62 A = numpy.identity(3, numpy.double)
  63 for i in range(len(files) - 1):
  64   in_jpg, out_jpg, top_left1 = files[i + 1]
  65   sys.stderr.write(f'read {in_jpg:s}\n')
  66   image1 = gamma.read_image(in_jpg)
  67   assert image1.shape == shape
  68
  69   p = []
  70   q = []
  71   for j in range(4):
  72     xs, ys = size[:, j]
  73     xf, yf = freedom[:, j]
  74     xf = int(round(xf * xs))
  75     yf = int(round(yf * ys))
  76     x0, y0 = top_left0[:, j]
  77     x1, y1 = top_left1[:, j]
  78
  79     block0 = image0[
  80       max(y0, 0):max(y0 + ys, 0),
  81       max(x0, 0):max(x0 + xs, 0),
  82       :
  83     ]
  84     #gamma.write_image(f'diag{j:d}.jpg', block0)
  85     mean0 = numpy.mean(numpy.mean(block0, 0), 0)
  86     data0 = block0 - mean0[numpy.newaxis, numpy.newaxis, :]
  87
  88     best = None
  89     for x in range(-xf, xf + 1):
  90       for y in range(-yf, yf + 1):
  91         block1 = image1[
  92           max(y1 + y, 0):max(y1 + y + ys, 0),
  93           max(x1 + x, 0):max(x1 + x + xs, 0),
  94           :
  95         ]
  96         if block1.shape != block0.shape:
  97           continue
  98         mean1 = numpy.mean(numpy.mean(block1, 0), 0)
  99         data1 = block1 - mean1[numpy.newaxis, numpy.newaxis, :]
 100
 101         score = numpy.sum(data0 * data1)
 102         if best is None or score > best[0]:
 103           best = (score, x, y)
 104
 105     score, x, y = best
 106     sys.stderr.write(f'score {score:.3f} x {x / xs:.3f} y {y / ys:.3f}\n')
 107     p.append(numpy.array([x0 + xs / 2, y0 + ys / 2], numpy.double))
 108     q.append(numpy.array([x1 + x + xs / 2, y1 + y + ys / 2], numpy.double))
 109
 110     top_left1[:, j] = (x1 + x, y1 + y)
 111   p = numpy.stack(p, 1)
 112   q = numpy.stack(q, 1)
 113   A = perspective.calc_transform(p, q) @ A # makes the transform cumulative
 114
 115   sys.stderr.write('remap\n')
 116   ys, xs, cs = shape
 117   #def mapping(p):
 118   #  #return tuple(
 119   #  #  perspective.apply_transform(A, numpy.array(p, numpy.double)[::-1])[::-1]
 120   #  #)
 121   #  q = A[:, 0] * p[1] + A[:, 1] * p[0] + A[:, 2]
 122   #  return (q[1] / q[2], q[0] / q[2])
 123   #out_image1 = numpy.stack(
 124   #  [
 125   #    scipy.ndimage.geometric_transform(image1[:, :, j], mapping)
 126   #    for j in range(cs)
 127   #  ],
 128   #  2
 129   #)
 130   coords = numpy.zeros((2, ys, xs), numpy.double)
 131   coords[0, :, :] = numpy.arange(ys, dtype = numpy.double)[:, numpy.newaxis]
 132   coords[1, :, :] = numpy.arange(xs, dtype = numpy.double)[numpy.newaxis, :]
 133   mapped_coords = perspective.apply_transform_multi(
 134     A,
 135     coords[::-1, :, :].reshape((2, ys * xs))
 136   ).reshape(2, ys, xs)[::-1, :, :]
 137   out_image1 = numpy.stack(
 138     [
 139       scipy.ndimage.map_coordinates(image1[:, :, j], mapped_coords)
 140       for j in range(cs)
 141     ],
 142     2
 143   )
 144
 145   sys.stderr.write(f'write {out_jpg:s}\n')
 146   gamma.write_image(out_jpg, out_image1)
 147
 148   image0 = image1
 149   top_left0 = top_left1
 150
 151 sys.stdout.write(
 152   ' '.join(
 153     [
 154       ','.join([str(size[j, i]) for j in range(2)])
 155       for i in range(4)
 156     ]
 157   ) + '\n'
 158 )
 159 sys.stdout.write(
 160   ' '.join(
 161     [
 162       ','.join([str(freedom[j, i]) for j in range(2)])
 163       for i in range(4)
 164     ]
 165   ) + '\n'
 166 )
 167 for in_jpg, out_jpg, top_left in files:
 168   sys.stdout.write(
 169     ' '.join(
 170       [in_jpg, out_jpg] +
 171         [
 172           ','.join([str(top_left[j, i]) for j in range(2)])
 173           for i in range(4)
 174         ]
 175     ) + '\n'
 176   )