manual.py

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