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     corners = []
  31     color=(64, 64, 255)
  32     line_width = 2
  33     while not done:
  34         for event in pygame.event.get():
  35             if event.type == pygame.QUIT:
  36                 pygame.quit()
  37                 raise UserQuitError
  38             if event.type == pygame.MOUSEBUTTONDOWN:
  39                if len(corners) >= 4:
  40                     corners = []
  41                     im = im_orig.copy()
  42                     draw = ImageDraw.Draw(im)
  43
  44                if len(corners) < 4:
  45                     corners.append(pygame.mouse.get_pos())
  46                     draw.point(corners[:-1], fill=color)
  47                     if len(corners) == 4:
  48                         draw.line((corners[0], corners[1]), fill=color,
  49                                   width=line_width)
  50                         draw.line((corners[1], corners[2]), fill=color,
  51                                   width=line_width)
  52                         draw.line((corners[2], corners[3]), fill=color,
  53                                   width=line_width)
  54                         draw.line((corners[3], corners[0]), fill=color,
  55                                   width=line_width)
  56                         l1 = half_line(corners)
  57                         draw.line(l1, fill=color, width=line_width)
  58                         l2 = half_line(corners[1:4] + [corners[0]])
  59                         draw.line(l2, fill=color, width=line_width)
  60                         c = center(corners)
  61                         draw.ellipse((c[0] - 1, c[1] - 1, c[0] + 1, c[1] + 1),
  62                                      fill=(255, 64, 64))
  63
  64         screen.display_picture(im)
  65         clock.tick(15)
  66
  67 def half_line(corners):
  68     c = center(corners)
  69     d1 = intersection(line(corners[0], corners[3]),
  70                       line(corners[1], corners[2]))
  71     p1 = intersection(line(c,d1), line(corners[0],
  72                                       corners[1]))
  73     p2 = intersection(line(c,d1), line(corners[2],
  74                                       corners[3]))
  75     return (p1, p2)
  76
  77 def center(corners):
  78     return intersection(line(corners[0], corners[2]), line(corners[1],
  79                                                            corners[3]))
  80 def line(x, y):
  81     a = float(x[1] - y[1])
  82     b = float(y[0] - x[0])
  83     c = a * y[0] + b * y[1]
  84     return (a, b, c)
  85
  86 def intersection(p, q):
  87     det = p[0] * q[1] - p[1] * q[0]
  88     if det == 0:
  89         return None
  90     return (int(round((q[1] * p[2] - p[1] * q[2]) / det)), int(round((p[0] *
  91                                                 q[2] - q[0] * p[2]) / det)))