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 find_lines(im_orig):
  29     # TODO rename, refactor, comment
  30
  31     im = im_orig.copy()
  32
  33     screen = Screen(im.size)
  34
  35     clock = pygame.time.Clock()
  36     draw = ImageDraw.Draw(im)
  37     hoshi = lambda c: draw.ellipse((c[0] - 1, c[1] - 1, c[0] + 1, c[1] + 1),
  38                  fill=(255, 64, 64))
  39     corners = []
  40     color = (32, 255, 32)
  41     line_width = 1
  42     lines_r = []
  43
  44     while True:
  45         for event in pygame.event.get():
  46             if event.type == pygame.QUIT or event.type == pygame.KEYDOWN:
  47                 pygame.quit()
  48                 if len(corners) == 4:
  49                     return lines_r
  50                 else:
  51                     raise UserQuitError
  52             if event.type == pygame.MOUSEBUTTONDOWN:
  53                 np = pygame.mouse.get_pos()
  54                 if len(corners) >= 4:
  55                     corners.sort(key=lambda p: dst(p, np))
  56                     corners = corners[1:]
  57                 corners.append(np)
  58                 (x, y) = corners[-1]
  59                 draw.line((x-2, y, x + 2, y), fill=color)
  60                 draw.line((x, y+2, x, y-2), fill=color)
  61                 if len(corners) == 4:
  62                     im = im_orig.copy()
  63                     draw = ImageDraw.Draw(im)
  64                     l_vert, l_hor = lines(corners)
  65                     for l in l_vert:
  66                         draw.line(l, fill=color, width=line_width)
  67                     for l in l_hor:
  68                         draw.line(l, fill=color, width=line_width)
  69                     # TODO sort by distance
  70                     l_vert.sort()
  71                     l_hor.sort()
  72                     for i in [3, 9, 15]:
  73                         for j in [3, 9, 15]:
  74                             hoshi(intersection(line(l_vert[i][0], l_vert[i][1]),
  75                                                line(l_hor[j][0], l_hor[j][1])))
  76                     lines_r = [[l2ad(l, im.size) for l in l_vert],
  77                                [l2ad(l, im.size) for l in l_hor]]
  78
  79         screen.display_picture(im)
  80         clock.tick(15)
  81
  82 def lines(corners):
  83     # TODO Error on triangle
  84     corners.sort() # TODO does this help?
  85     # TODO refactor this vvv
  86     cor_d = [(corners[0], (c[0] - corners[0][0], c[1] - corners[0][1]), c) for c in
  87              corners[1:]]
  88     cor_d = [(float(a[0] * b[0] + a[1] * b[1]) / (sqrt(a[0] ** 2 + a[1] ** 2) *
  89               sqrt(b[0] **2 + b[1] ** 2)), a[0] * b[1] - b[0] * a[1], c) for a, b, c in cor_d]
  90     cor_d = sorted([(copysign(acos(min(a, 1)), b), c) for a, b, c in cor_d])
  91     corners = [corners[0]] + [c for _, c in cor_d]
  92     return (_lines(corners, 0) +
  93             [(corners[0], corners[3]), (corners[1], corners[2])],
  94             _lines(corners[1:4] + [corners[0]], 0) +
  95             [(corners[0], corners[1]), (corners[2], corners[3])])
  96
  97 def _lines(corners, n):
  98     # TODO what is this?
  99     if n == 0:
 100         x = half_line(corners)
 101         return (_lines([corners[0], x[0], x[1], corners[3]], 1) + [x] +
 102                 _lines([x[0], corners[1], corners[2], x[1]], 1))
 103     else:
 104         x = half_line(corners)
 105         c = intersection(line(x[0], corners[2]), line(corners[1], corners[3]))
 106         d = intersection(line(corners[0], corners[3]), line(corners[1], corners[2]))
 107         if d:
 108             l = (intersection(line(corners[0], corners[1]), line(c, d)),
 109                  intersection(line(corners[2], corners[3]), line(c, d)))
 110         else:
 111             lx = line(c, (c[0] + corners[0][0] - corners[3][0],
 112                       c[1] + corners[0][1] - corners[3][1]))
 113             l = (intersection(line(corners[0], corners[1]), lx),
 114                  intersection(line(corners[2], corners[3]), lx))
 115         l2 = half_line([corners[0], l[0], l[1], corners[3]])
 116         if n == 1:
 117             return ([l, l2] + _lines([l[0], l2[0], l2[1], l[1]], 2)
 118                     + _lines([corners[0], l2[0], l2[1], corners[3]], 2)
 119                     + _lines([l[0], corners[1], corners[2], l[1]], 2))
 120         if n == 2:
 121             return [l, l2]
 122
 123
 124 def half_line(corners):
 125     # TODO what is this?
 126     c = center(corners)
 127     d = intersection(line(corners[0], corners[3]), line(corners[1], corners[2]))
 128     if d:
 129         l = line(c, d)
 130     else:
 131         l = line(c, (c[0] + corners[0][0] - corners[3][0],
 132                      c[1] + corners[0][1] - corners[3][1]))
 133     p1 = intersection(l, line(corners[0], corners[1]))
 134     p2 = intersection(l, line(corners[2], corners[3]))
 135     return (p1, p2)
 136
 137 def center(corners):
 138     """Given a list of four corner points, return the center of the square."""
 139     return intersection(line(corners[0], corners[2]),
 140                         line(corners[1], corners[3]))