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