imago_pack/manual.py

   1 """Manual grid selection module"""
   2
   3 import pygame
   4 import ImageDraw
   5 from math import sqrt, acos, copysign
   6
   7 from geometry import l2ad, line, intersection
   8
   9 class UserQuitError(Exception):
  10     pass
  11
  12 class Screen:
  13     def __init__(self, res):
  14         pygame.init()
  15         pygame.display.set_mode(res)
  16         pygame.display.set_caption("Imago manual mode")
  17         self._screen = pygame.display.get_surface()
  18
  19     def display_picture(self, img):
  20         pg_img = pygame.image.frombuffer(img.tostring(), img.size, img.mode)
  21         self._screen.blit(pg_img, (0, 0))
  22         pygame.display.flip()
  23
  24 def find_lines(im_orig):
  25
  26     im = im_orig.copy()
  27
  28     screen = Screen(im.size)
  29
  30     done = False
  31     clock = pygame.time.Clock()
  32     draw = ImageDraw.Draw(im)
  33     hoshi = lambda c: draw.ellipse((c[0] - 1, c[1] - 1, c[0] + 1, c[1] + 1),
  34                  fill=(255, 64, 64))
  35     corners = []
  36     color = (64, 64, 255)
  37     line_width = 1
  38     lines_r = []
  39
  40     while not done:
  41         for event in pygame.event.get():
  42             if event.type == pygame.QUIT or event.type == pygame.KEYDOWN:
  43                 pygame.quit()
  44                 if len(corners) == 4:
  45                     return lines_r
  46                 else:
  47                     raise UserQuitError
  48             if event.type == pygame.MOUSEBUTTONDOWN:
  49                 if len(corners) >= 4:
  50                     corners = []
  51                     im = im_orig.copy()
  52                     draw = ImageDraw.Draw(im)
  53
  54                 if len(corners) < 4:
  55                     corners.append(pygame.mouse.get_pos())
  56                     draw.point(corners[:-1], fill=color)
  57                     if len(corners) == 4:
  58                         l_vert, l_hor = lines(corners)
  59                         for l in l_vert:
  60                             draw.line(l, fill=color, width=line_width)
  61                         for l in l_hor:
  62                             draw.line(l, fill=color, width=line_width)
  63                         #TODO sort by distance
  64                         l_vert.sort()
  65                         l_hor.sort()
  66                         for i in [3, 9, 15]:
  67                             for j in [3, 9, 15]:
  68                                 hoshi(intersection(line(l_vert[i][0], l_vert[i][1]),
  69                                                    line(l_hor[j][0], l_hor[j][1])))
  70                         lines_r = [[l2ad(l, im.size) for l in l_vert],
  71                                    [l2ad(l, im.size) for l in l_hor]]
  72
  73         screen.display_picture(im)
  74         clock.tick(15)
  75
  76 def lines(corners):
  77     #TODO Error on triangle
  78     cor_d = [(corners[0], (c[0] - corners[0][0], c[1] - corners[0][1]), c) for c in
  79              corners[1:]]
  80     cor_d = [(float(a[0] * b[0] + a[1] * b[1]) / (sqrt(a[0] ** 2 + a[1] ** 2) *
  81               sqrt(b[0] **2 + b[1] ** 2)), a[0] * b[1] - b[0] * a[1], c) for a, b, c in cor_d]
  82     cor_d = sorted([(copysign(acos(min(a, 1)), b), c) for a, b, c in cor_d])
  83     corners = [corners[0]] + [c for _, c in cor_d]
  84     return (_lines(corners, 0) + [(corners[0], corners[3]),
  85                                   (corners[1], corners[2])],
  86             _lines(corners[1:4] + [corners[0]], 0) +
  87             [(corners[0], corners[1]), (corners[2], corners[3])])
  88
  89 def _lines(corners, n):
  90     if n == 0:
  91         x = half_line(corners)
  92         return (_lines([corners[0], x[0], x[1], corners[3]], n + 1) + [x] +
  93                 _lines([x[0], corners[1], corners[2], x[1]], n + 1))
  94     else:
  95         x = half_line(corners)
  96         c = intersection(line(x[0], corners[2]), line(corners[1], corners[3]))
  97         d = intersection(line(corners[0], corners[3]), line(corners[1], corners[2]))
  98         if d:
  99             l = (intersection(line(corners[0], corners[1]), line(c, d)),
 100                  intersection(line(corners[2], corners[3]), line(c, d)))
 101         else:
 102             lx = line(c, (c[0] + corners[0][0] - corners[3][0],
 103                       c[1] + corners[0][1] - corners[3][1]))
 104             l = (intersection(line(corners[0], corners[1]), lx),
 105                  intersection(line(corners[2], corners[3]), lx))
 106         l2 = half_line([corners[0], l[0], l[1], corners[3]])
 107         if n == 1:
 108             return ([l, l2] + _lines([l[0], l2[0], l2[1], l[1]], 2)
 109                     + _lines([corners[0], l2[0], l2[1], corners[3]], 2)
 110                     + _lines([l[0], corners[1], corners[2], l[1]], 2))
 111         if n == 2:
 112             return [l, l2]
 113
 114
 115 def half_line(corners):
 116     c = center(corners)
 117     d = intersection(line(corners[0], corners[3]), line(corners[1], corners[2]))
 118     if d:
 119         l = line(c, d)
 120     else:
 121         l = line(c, (c[0] + corners[0][0] - corners[3][0],
 122                      c[1] + corners[0][1] - corners[3][1]))
 123     p1 = intersection(l, line(corners[0], corners[1]))
 124     p2 = intersection(l, line(corners[2], corners[3]))
 125     return (p1, p2)
 126
 127
 128 def center(corners):
 129     return intersection(line(corners[0], corners[2]),
 130                         line(corners[1], corners[3]))