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 from manual_lines import *
   9
  10 class UserQuitError(Exception):
  11     pass
  12
  13 class Screen:
  14     # TODO isn't this a duplicate of something?
  15     def __init__(self, res):
  16         pygame.init()
  17         pygame.display.set_mode(res)
  18         pygame.display.set_caption("Imago manual mode")
  19         self._screen = pygame.display.get_surface()
  20
  21     def display_picture(self, img):
  22         pg_img = pygame.image.frombuffer(img.tostring(), img.size, img.mode)
  23         self._screen.blit(pg_img, (0, 0))
  24         pygame.display.flip()
  25
  26
  27 def display_instr():
  28
  29     while True:
  30         for event in pygame.event.get():
  31             if event.type == pygame.QUIT:
  32                 pygame.quit()
  33                 raise UserQuitError
  34             if event.type == pygame.MOUSEBUTTONDOWN or event.type == pygame.KEYDOWN:
  35                 return
  36
  37 def find_lines(im_orig):
  38     # TODO rename, refactor, comment
  39
  40     im = im_orig.copy()
  41
  42     screen = Screen((620, 350))
  43
  44     font = pygame.font.Font(None, 25)
  45     instructions = ["Imago manual mode", "",
  46     "To set the grid position, click on the corners of the grid. Once you mark",
  47     "all four corners, the grid will appear. To adjust it, just click on the new",
  48     "position and the nearest corner will move there. Once you are content",
  49     "with the alignment, press any key on your keyboard or close the window.",
  50     "", "", "",
  51     "Press any key to continue."]
  52     y = 10
  53     for i in instructions:
  54         text1 = font.render(i, True, [128, 255, 128])
  55         screen._screen.blit(text1, [10, y])
  56         y += 25
  57
  58     pygame.display.flip()
  59
  60     display_instr()
  61
  62     pygame.display.set_mode(im.size)
  63
  64     clock = pygame.time.Clock()
  65     draw = ImageDraw.Draw(im)
  66     hoshi = lambda c: draw.ellipse((c[0] - 1, c[1] - 1, c[0] + 1, c[1] + 1),
  67                  fill=(255, 64, 64))
  68     corners = []
  69     color = (32, 255, 32)
  70     line_width = 1
  71     lines_r = []
  72
  73     def dst((x1, y1), (x2, y2)):
  74         return (x1 - x2) ** 2 + (y1 - y2) ** 2
  75
  76     while True:
  77         for event in pygame.event.get():
  78             if event.type == pygame.QUIT or event.type == pygame.KEYDOWN:
  79                 pygame.quit()
  80                 if len(corners) == 4:
  81                     return lines_r
  82                 else:
  83                     raise UserQuitError
  84             if event.type == pygame.MOUSEBUTTONDOWN:
  85                 np = pygame.mouse.get_pos()
  86                 if len(corners) >= 4:
  87                     corners.sort(key=lambda p: dst(p, np))
  88                     corners = corners[1:]
  89                 corners.append(np)
  90                 (x, y) = corners[-1]
  91                 draw.line((x-2, y, x + 2, y), fill=color)
  92                 draw.line((x, y+2, x, y-2), fill=color)
  93                 if len(corners) == 4:
  94                     im = im_orig.copy()
  95                     draw = ImageDraw.Draw(im)
  96                     try:
  97                         l_vert, l_hor = lines(corners)
  98                     except Exception:
  99                         corners = corners[:-1]
 100                     for l in l_vert:
 101                         draw.line(l, fill=color, width=line_width)
 102                     for l in l_hor:
 103                         draw.line(l, fill=color, width=line_width)
 104                     # TODO sort by distance
 105                     #l_vert.sort()
 106                     #l_hor.sort()
 107                     #for i in [3, 9, 15]:
 108                     #    for j in [3, 9, 15]:
 109                     #        hoshi(intersection(line(l_vert[i][0], l_vert[i][1]),
 110                     #                           line(l_hor[j][0], l_hor[j][1])))
 111                     lines_r = [[l2ad(l, im.size) for l in l_vert],
 112                                [l2ad(l, im.size) for l in l_hor]]
 113
 114         screen.display_picture(im)
 115         clock.tick(15)
 116
 117