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                     try:
  95                         l_vert, l_hor = lines(corners)
  96                     except Exception:
  97                         corners = corners[:-1]
  98                     for l in l_vert:
  99                         draw.line(l, fill=color, width=line_width)
 100                     for l in l_hor:
 101                         draw.line(l, fill=color, width=line_width)
 102                     # TODO sort by distance
 103                     #l_vert.sort()
 104                     #l_hor.sort()
 105                     #for i in [3, 9, 15]:
 106                     #    for j in [3, 9, 15]:
 107                     #        hoshi(intersection(line(l_vert[i][0], l_vert[i][1]),
 108                     #                           line(l_hor[j][0], l_hor[j][1])))
 109                     lines_r = [[l2ad(l, im.size) for l in l_vert],
 110                                [l2ad(l, im.size) for l in l_hor]]
 111
 112         screen.display_picture(im)
 113         clock.tick(15)
 114
 115 def lines(corners):
 116     # TODO Error on triangle
 117     corners.sort() # TODO does this help?
 118     # TODO refactor this vvv
 119     cor_d = [(corners[0], (c[0] - corners[0][0], c[1] - corners[0][1]), c) for c in
 120              corners[1:]]
 121     cor_d = [(float(a[0] * b[0] + a[1] * b[1]) / (sqrt(a[0] ** 2 + a[1] ** 2) *
 122               sqrt(b[0] **2 + b[1] ** 2)), a[0] * b[1] - b[0] * a[1], c) for a, b, c in cor_d]
 123     cor_d = sorted([(copysign(acos(min(a, 1)), b), c) for a, b, c in cor_d])
 124     corners = [corners[0]] + [c for _, c in cor_d]
 125     return (_lines(corners, 0) +
 126             [(corners[0], corners[3]), (corners[1], corners[2])],
 127             _lines(corners[1:4] + [corners[0]], 0) +
 128             [(corners[0], corners[1]), (corners[2], corners[3])])
 129
 130 def _lines(corners, n):
 131     # TODO what is this?
 132     if n == 0:
 133         x = half_line(corners)
 134         return (_lines([corners[0], x[0], x[1], corners[3]], 1) + [x] +
 135                 _lines([x[0], corners[1], corners[2], x[1]], 1))
 136     else:
 137         x = half_line(corners)
 138         c = intersection(line(x[0], corners[2]), line(corners[1], corners[3]))
 139         d = intersection(line(corners[0], corners[3]), line(corners[1], corners[2]))
 140         if d:
 141             l = (intersection(line(corners[0], corners[1]), line(c, d)),
 142                  intersection(line(corners[2], corners[3]), line(c, d)))
 143         else:
 144             lx = line(c, (c[0] + corners[0][0] - corners[3][0],
 145                       c[1] + corners[0][1] - corners[3][1]))
 146             l = (intersection(line(corners[0], corners[1]), lx),
 147                  intersection(line(corners[2], corners[3]), lx))
 148         l2 = half_line([corners[0], l[0], l[1], corners[3]])
 149         if n == 1:
 150             return ([l, l2] + _lines([l[0], l2[0], l2[1], l[1]], 2)
 151                     + _lines([corners[0], l2[0], l2[1], corners[3]], 2)
 152                     + _lines([l[0], corners[1], corners[2], l[1]], 2))
 153         if n == 2:
 154             return [l, l2]
 155
 156
 157 def half_line(corners):
 158     # TODO what is this?
 159     c = center(corners)
 160     d = intersection(line(corners[0], corners[3]), line(corners[1], corners[2]))
 161     if d:
 162         l = line(c, d)
 163     else:
 164         l = line(c, (c[0] + corners[0][0] - corners[3][0],
 165                      c[1] + corners[0][1] - corners[3][1]))
 166     p1 = intersection(l, line(corners[0], corners[1]))
 167     p2 = intersection(l, line(corners[2], corners[3]))
 168     return (p1, p2)
 169
 170 def center(corners):
 171     """Given a list of four corner points, return the center of the square."""
 172     return intersection(line(corners[0], corners[2]),
 173                         line(corners[1], corners[3]))