src/manual.py

   1 """Manual grid selection module"""
   2
   3 from math import sqrt, acos, copysign
   4 from PIL import ImageDraw
   5 import pygame
   6
   7 from geometry import l2ad, line, intersection
   8
   9 class UserQuitError(Exception):
  10     pass
  11
  12 class Screen:
  13     # TODO isn't this a duplicate of something?
  14     def __init__(self, res):
  15         pygame.init()
  16         pygame.display.set_mode(res)
  17         pygame.display.set_caption("Imago manual mode")
  18         self._screen = pygame.display.get_surface()
  19
  20     def display_picture(self, img):
  21         pg_img = pygame.image.frombuffer(img.tostring(), img.size, img.mode)
  22         self._screen.blit(pg_img, (0, 0))
  23         pygame.display.flip()
  24
  25 def dst((x1, y1), (x2, y2)):
  26     return (x1 - x2) ** 2 + (y1 - y2) ** 2
  27
  28 def display_instr():
  29
  30     while True:
  31         for event in pygame.event.get():
  32             if event.type == pygame.QUIT:
  33                 pygame.quit()
  34                 raise UserQuitError
  35             if event.type == pygame.MOUSEBUTTONDOWN or event.type == pygame.KEYDOWN:
  36                 return
  37
  38 def find_lines(im_orig):
  39     # TODO rename, refactor, comment
  40
  41     im = im_orig.copy()
  42
  43     screen = Screen((620, 350))
  44
  45     font = pygame.font.Font(None, 25)
  46     instructions = ["Imago manual mode", "",
  47     "To set the grid position, click on the corners of the grid. Once you mark",
  48     "all four corners, the grid will appear. To adjust it, just click on the new",
  49     "position and the nearest corner will move there. Once you are content",
  50     "with the alignment, press any key on your keyboard or close the window.",
  51     "", "", "",
  52     "Press any key to continue."]
  53     y = 10
  54     for i in instructions:
  55         text1 = font.render(i, True, [128, 255, 128])
  56         screen._screen.blit(text1, [10, y])
  57         y += 25
  58
  59     pygame.display.flip()
  60
  61     display_instr()
  62
  63     pygame.display.set_mode(im.size)
  64
  65     clock = pygame.time.Clock()
  66     draw = ImageDraw.Draw(im)
  67     hoshi = lambda c: draw.ellipse((c[0] - 1, c[1] - 1, c[0] + 1, c[1] + 1),
  68                  fill=(255, 64, 64))
  69     corners = []
  70     color = (32, 255, 32)
  71     line_width = 1
  72     lines_r = []
  73
  74     while True:
  75         for event in pygame.event.get():
  76             if event.type == pygame.QUIT or event.type == pygame.KEYDOWN:
  77                 pygame.quit()
  78                 if len(corners) == 4:
  79                     return lines_r
  80                 else:
  81                     raise UserQuitError
  82             if event.type == pygame.MOUSEBUTTONDOWN:
  83                 np = pygame.mouse.get_pos()
  84                 if len(corners) >= 4:
  85                     corners.sort(key=lambda p: dst(p, np))
  86                     corners = corners[1:]
  87                 corners.append(np)
  88                 (x, y) = corners[-1]
  89                 draw.line((x-2, y, x + 2, y), fill=color)
  90                 draw.line((x, y+2, x, y-2), fill=color)
  91                 if len(corners) == 4:
  92                     im = im_orig.copy()
  93                     draw = ImageDraw.Draw(im)
  94                     l_vert, l_hor = lines(corners)
  95                     for l in l_vert:
  96                         draw.line(l, fill=color, width=line_width)
  97                     for l in l_hor:
  98                         draw.line(l, fill=color, width=line_width)
  99                     # TODO sort by distance
 100                     #l_vert.sort()
 101                     #l_hor.sort()
 102                     #for i in [3, 9, 15]:
 103                     #    for j in [3, 9, 15]:
 104                     #        hoshi(intersection(line(l_vert[i][0], l_vert[i][1]),
 105                     #                           line(l_hor[j][0], l_hor[j][1])))
 106                     lines_r = [[l2ad(l, im.size) for l in l_vert],
 107                                [l2ad(l, im.size) for l in l_hor]]
 108
 109         screen.display_picture(im)
 110         clock.tick(15)
 111
 112 def lines(corners):
 113     # TODO Error on triangle
 114     corners.sort() # TODO does this help?
 115     # TODO refactor this vvv
 116     cor_d = [(corners[0], (c[0] - corners[0][0], c[1] - corners[0][1]), c) for c in
 117              corners[1:]]
 118     cor_d = [(float(a[0] * b[0] + a[1] * b[1]) / (sqrt(a[0] ** 2 + a[1] ** 2) *
 119               sqrt(b[0] **2 + b[1] ** 2)), a[0] * b[1] - b[0] * a[1], c) for a, b, c in cor_d]
 120     cor_d = sorted([(copysign(acos(min(a, 1)), b), c) for a, b, c in cor_d])
 121     corners = [corners[0]] + [c for _, c in cor_d]
 122     return (_lines(corners, 0) +
 123             [(corners[0], corners[3]), (corners[1], corners[2])],
 124             _lines(corners[1:4] + [corners[0]], 0) +
 125             [(corners[0], corners[1]), (corners[2], corners[3])])
 126
 127 def _lines(corners, n):
 128     # TODO what is this?
 129     if n == 0:
 130         x = half_line(corners)
 131         return (_lines([corners[0], x[0], x[1], corners[3]], 1) + [x] +
 132                 _lines([x[0], corners[1], corners[2], x[1]], 1))
 133     else:
 134         x = half_line(corners)
 135         c = intersection(line(x[0], corners[2]), line(corners[1], corners[3]))
 136         d = intersection(line(corners[0], corners[3]), line(corners[1], corners[2]))
 137         if d:
 138             l = (intersection(line(corners[0], corners[1]), line(c, d)),
 139                  intersection(line(corners[2], corners[3]), line(c, d)))
 140         else:
 141             lx = line(c, (c[0] + corners[0][0] - corners[3][0],
 142                       c[1] + corners[0][1] - corners[3][1]))
 143             l = (intersection(line(corners[0], corners[1]), lx),
 144                  intersection(line(corners[2], corners[3]), lx))
 145         l2 = half_line([corners[0], l[0], l[1], corners[3]])
 146         if n == 1:
 147             return ([l, l2] + _lines([l[0], l2[0], l2[1], l[1]], 2)
 148                     + _lines([corners[0], l2[0], l2[1], corners[3]], 2)
 149                     + _lines([l[0], corners[1], corners[2], l[1]], 2))
 150         if n == 2:
 151             return [l, l2]
 152
 153
 154 def half_line(corners):
 155     # TODO what is this?
 156     c = center(corners)
 157     d = intersection(line(corners[0], corners[3]), line(corners[1], corners[2]))
 158     if d:
 159         l = line(c, d)
 160     else:
 161         l = line(c, (c[0] + corners[0][0] - corners[3][0],
 162                      c[1] + corners[0][1] - corners[3][1]))
 163     p1 = intersection(l, line(corners[0], corners[1]))
 164     p2 = intersection(l, line(corners[2], corners[3]))
 165     return (p1, p2)
 166
 167 def center(corners):
 168     """Given a list of four corner points, return the center of the square."""
 169     return intersection(line(corners[0], corners[2]),
 170                         line(corners[1], corners[3]))