src/intrsc.py

   1 """Imago intersections module."""
   2
   3 from math import cos, tan, pi
   4 from operator import itemgetter
   5
   6 import ImageDraw
   7
   8 import filters
   9 import k_means
  10 import output
  11
  12 def dst(line):
  13     """Return normalized line."""
  14     if line[0] < pi / 2:
  15         line = line[0] + pi, - line[1]
  16     return line
  17
  18 def dst_sort(lines):
  19     """Return lines sorted by distance."""
  20     l_max = max(l[0] for l in lines)
  21     l_min = min(l[0] for l in lines)
  22     if l_max - l_min > (3. / 4) * pi:
  23         lines = [dst(l) for l in lines]
  24     lines.sort(key=itemgetter(1))
  25     return lines
  26
  27 def board(image, lines, show_all, do_something):
  28     """Compute intersections, find stone colors and return board situation."""
  29     # TODO refactor show_all, do_something
  30     # TODO refactor this into smaller functions
  31     lines = [dst_sort(l) for l in lines]
  32     intersections = intersections_from_angl_dist(lines, image.size)
  33
  34     if show_all:
  35         image_g = image.copy()
  36         draw = ImageDraw.Draw(image_g)
  37         for line in intersections:
  38             for (x, y) in line:
  39                 draw.point((x , y), fill=(120, 255, 120))
  40         do_something(image_g, "intersections")
  41
  42     image_c = filters.color_enhance(image)
  43     if show_all:
  44         do_something(image_c, "white balance")
  45
  46     board_raw = []
  47
  48     for line in intersections:
  49         board_raw.append([stone_color_raw(image_c, intersection) for intersection in
  50                       line])
  51     board_raw = sum(board_raw, [])
  52
  53     ### Show color distribution
  54     luma = [s[0] for s in board_raw]
  55     saturation = [s[1] for s in board_raw]
  56
  57     if show_all:
  58         import matplotlib.pyplot as pyplot
  59         import Image
  60         fig = pyplot.figure(figsize=(8, 6))
  61         pyplot.scatter(luma, saturation,
  62                        color=[(s[2][0]/255.,
  63                                s[2][1]/255.,
  64                                s[2][2]/255., 1.)
  65                                    for s in board_raw])
  66         pyplot.xlim(0,1)
  67         pyplot.ylim(0,1)
  68         fig.canvas.draw()
  69         size = fig.canvas.get_width_height()
  70         buff = fig.canvas.tostring_rgb()
  71         image_p = Image.fromstring('RGB', size, buff, 'raw')
  72         do_something(image_p, "color distribution")
  73
  74     clusters = k_means.cluster(3, 2,zip(zip(luma, saturation), range(len(luma))),
  75                                [[0., 0.5], [0.5, 0.5], [1., 0.5]])
  76
  77     if show_all:
  78         fig = pyplot.figure(figsize=(8, 6))
  79         pyplot.scatter([d[0][0] for d in clusters[0]], [d[0][1] for d in clusters[0]],
  80                                                  color=(1,0,0,1))
  81         pyplot.scatter([d[0][0] for d in clusters[1]], [d[0][1] for d in clusters[1]],
  82                                                  color=(0,1,0,1))
  83         pyplot.scatter([d[0][0] for d in clusters[2]], [d[0][1] for d in clusters[2]],
  84                                                  color=(0,0,1,1))
  85         pyplot.xlim(0,1)
  86         pyplot.ylim(0,1)
  87         fig.canvas.draw()
  88         size = fig.canvas.get_width_height()
  89         buff = fig.canvas.tostring_rgb()
  90         image_p = Image.fromstring('RGB', size, buff, 'raw')
  91         do_something(image_p, "color clustering")
  92
  93     clusters[0] = [(p[1], 'B') for p in clusters[0]]
  94     clusters[1] = [(p[1], '.') for p in clusters[1]]
  95     clusters[2] = [(p[1], 'W') for p in clusters[2]]
  96
  97     board_rl = sum(clusters, [])
  98     board_rl.sort()
  99     board_rg = (p[1] for p in board_rl)
 100
 101     board_r = []
 102
 103     #TODO 19 should be a size parameter
 104     try:
 105         for i in xrange(19):
 106             for _ in xrange(19):
 107                 board_r.append(board_rg.next())
 108     except StopIteration:
 109         pass
 110
 111
 112     return output.Board(19, board_r)
 113
 114 def mean_luma(cluster):
 115     """Return mean luma of the *cluster* of points."""
 116     return sum(c[0][0] for c in cluster) / float(len(cluster))
 117
 118 def intersections_from_angl_dist(lines, size, get_all=True):
 119     """Take grid-lines and size of the image. Return intersections."""
 120     intersections = []
 121     for (angl1, dist1) in lines[1]:
 122         line = []
 123         for (angl2, dist2) in lines[0]:
 124             if abs(angl1 - angl2) > 0.4:
 125                 i_x =  (- ((dist2 / cos(angl2)) - (dist1 / cos(angl1)))
 126                         / (tan(angl1) - tan(angl2)))
 127                 i_y = (tan(angl1) * i_x) - (dist1 / cos(angl1))
 128                 if get_all or (-size[0] / 2 < i_x < size[0] / 2 and
 129                     -size[1] / 2 < i_y < size[1] / 2):
 130                     line.append((int(i_x + size[0] / 2),
 131                                  int(i_y + size[1] / 2)))
 132         intersections.append(line)
 133     return intersections
 134
 135 def rgb2lumsat(color):
 136     """Convert RGB to luma and HSI model saturation."""
 137     r, g, b = color
 138     luma = (0.30 * r + 0.59 * g + 0.11 * b) / 255.0
 139     max_diff = max(color) - min(color)
 140     if max_diff == 0:
 141         saturation = 0
 142     else:
 143         saturation = 1. - ((3. * min(color)) / sum(color))
 144     return luma, saturation
 145
 146 def median(lst):
 147     #TODO comment (or delete maybe?)
 148     len_lst = len(lst)
 149     if len_lst % 2 == 0:
 150         return (lst[len_lst / 2] + lst[len_lst / 2 + 1]) / 2.0
 151     else:
 152         return lst[len_lst / 2]
 153
 154 def stone_color_raw(image, (x, y)):
 155     """Given image and coordinates, return stone color."""
 156     size = 3
 157     points = []
 158     for i in range(-size, size + 1):
 159         for j in range(-size, size + 1):
 160             try:
 161                 points.append(image.getpixel((x + i, y + j)))
 162             except IndexError:
 163                 pass
 164     norm = float(len(points))
 165     color = (sum(p[0] for p in points) / norm,
 166              sum(p[1] for p in points) / norm,
 167              sum(p[2] for p in points) / norm)
 168     luma, saturation = rgb2lumsat(color)
 169     return luma, saturation, color