manual.py

   1 """Manual grid selection module"""
   2
   3 import pygame
   4 import Image, ImageDraw
   5 from math import atan, sin, cos, pi
   6
   7 class UserQuitError(Exception):
   8     pass
   9
  10 class Screen:
  11     def __init__(self, res):
  12         pygame.init()
  13         pygame.display.set_mode(res)
  14         pygame.display.set_caption("Imago manual mode")
  15         self._screen = pygame.display.get_surface()
  16
  17     def display_picture(self, im):
  18         pg_img = pygame.image.frombuffer(im.tostring(), im.size, im.mode)
  19         self._screen.blit(pg_img, (0,0))
  20         pygame.display.flip()
  21
  22 def find_lines(im_orig):
  23
  24     im = im_orig.copy()
  25
  26     screen = Screen(im.size)
  27
  28     done = False
  29     clock = pygame.time.Clock()
  30     draw = ImageDraw.Draw(im)
  31     hoshi = lambda c: draw.ellipse((c[0] - 1, c[1] - 1, c[0] + 1, c[1] + 1),
  32                  fill=(255, 64, 64))
  33     corners = []
  34     color=(64, 64, 255)
  35     line_width = 1
  36     lines_r = []
  37
  38     while not done:
  39         for event in pygame.event.get():
  40             if event.type == pygame.QUIT or event.type == pygame.KEYDOWN:
  41                 pygame.quit()
  42                 if len(corners) == 4:
  43                     return lines_r
  44                 else:
  45                     raise UserQuitError
  46             if event.type == pygame.MOUSEBUTTONDOWN:
  47                 if len(corners) >= 4:
  48                     corners = []
  49                     im = im_orig.copy()
  50                     draw = ImageDraw.Draw(im)
  51
  52                 if len(corners) < 4:
  53                     corners.append(pygame.mouse.get_pos())
  54                     draw.point(corners[:-1], fill=color)
  55                     if len(corners) == 4:
  56                         draw.line((corners[0], corners[1]), fill=color,
  57                                   width=line_width)
  58                         draw.line((corners[1], corners[2]), fill=color,
  59                                   width=line_width)
  60                         draw.line((corners[2], corners[3]), fill=color,
  61                                   width=line_width)
  62                         draw.line((corners[3], corners[0]), fill=color,
  63                                   width=line_width)
  64                         l_vert = lines(corners)
  65                         for l in l_vert:
  66                             draw.line(l, fill=color, width=line_width)
  67                         l_hor = lines(corners[1:4] + [corners[0]])
  68                         for l in l_hor:
  69                             draw.line(l, fill=color, width=line_width)
  70                         l_vert += [(corners[0], corners[3]),
  71                                    (corners[1], corners[2])]
  72                         l_hor += [(corners[0], corners[1]),
  73                                    (corners[2], corners[3])]
  74                         l_vert.sort()
  75                         l_hor.sort()
  76                         for i in [3, 9, 15]:
  77                             for j in [3, 9, 15]:
  78                                 hoshi(intersection(line(l_vert[i][0], l_vert[i][1]),
  79                                                    line(l_hor[j][0], l_hor[j][1])))
  80                         lines_r = [[l2ad(l[0], l[1], im.size) for l in l_vert],
  81                                    [l2ad(l[0], l[1], im.size) for l in l_hor]]
  82
  83         screen.display_picture(im)
  84         clock.tick(15)
  85
  86 def lines(corners):
  87     return _lines(corners, 0)
  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         l = (intersection(line(corners[0], corners[1]), line(c,d)),
  99              intersection(line(corners[2], corners[3]), line(c,d)))
 100         l2 = half_line([corners[0], l[0], l[1], corners[3]])
 101         if n == 1:
 102             return ([l, l2] + _lines([l[0], l2[0], l2[1], l[1]], 2)
 103                     + _lines([corners[0], l2[0], l2[1], corners[3]], 2)
 104                     + _lines([l[0], corners[1], corners[2], l[1]], 2))
 105         if n == 2:
 106             return [l, l2]
 107
 108
 109 def half_line(corners):
 110     c = center(corners)
 111     d = intersection(line(corners[0], corners[3]), line(corners[1], corners[2]))
 112     if d:
 113         l = line(c, d)
 114     else:
 115         l = line(c, (c[0] + corners[0][0] - corners[3][0],
 116                      c[1] + corners[0][1] - corners[3][1]))
 117     p1 = intersection(l, line(corners[0], corners[1]))
 118     p2 = intersection(l, line(corners[2], corners[3]))
 119     return (p1, p2)
 120
 121
 122 def center(corners):
 123     return intersection(line(corners[0], corners[2]),
 124                         line(corners[1], corners[3]))
 125 def line(x, y):
 126     a = x[1] - y[1]
 127     b = y[0] - x[0]
 128     c = a * y[0] + b * y[1]
 129     return (a, b, c)
 130
 131 def intersection(p, q):
 132     det = p[0] * q[1] - p[1] * q[0]
 133     if det == 0:
 134         return None
 135     return (int(round(float(q[1] * p[2] - p[1] * q[2]) / det)),
 136             int(round(float(p[0] * q[2] - q[0] * p[2]) / det)))
 137
 138 def l2ad(a, b, size):
 139     if (a[0] - b[0]) == 0:
 140         angle = pi / 2
 141     else:
 142         q = float(a[1] - b[1]) / (a[0] - b[0])
 143         angle = atan(q)
 144
 145     if angle < 0:
 146         angle += pi
 147     if angle > pi:
 148         angle -= pi
 149
 150     distance = (((a[0] - (size[0] / 2)) * sin(angle)) +
 151                 ((a[1] - (size[1] / 2)) * - cos(angle)))
 152     return (angle, distance)