src/intrsc.py

   1 """Imago intersections module."""
   2
   3 from math import cos, tan, pi
   4 from operator import itemgetter
   5 import colorsys
   6
   7 import ImageDraw
   8
   9 import filters
  10 import k_means
  11 import output
  12 import linef
  13
  14 def dst(line):
  15     """Return normalized line."""
  16     if line[0] < pi / 2:
  17         line = line[0] + pi, - line[1]
  18     return line
  19
  20 def dst_sort(lines):
  21     """Return lines sorted by distance."""
  22     l_max = max(l[0] for l in lines)
  23     l_min = min(l[0] for l in lines)
  24     if l_max - l_min > (3. / 4) * pi:
  25         lines = [dst(l) for l in lines]
  26     lines.sort(key=itemgetter(1))
  27     return lines
  28
  29 def b_intersects(image, lines, show_all, do_something, logger):
  30     """Compute intersections."""
  31     # TODO refactor show_all, do_something
  32     # TODO refactor this into smaller functions
  33     logger("finding the stones")
  34     lines = [dst_sort(l) for l in lines]
  35     an0 = (sum([l[0] for l in lines[0]]) / len(lines[0]) - pi / 2)
  36     an1 = (sum([l[0] for l in lines[1]]) / len(lines[1]) - pi / 2)
  37     if an0 > an1:
  38         lines = [lines[1], lines[0]]
  39
  40     intersections = intersections_from_angl_dist(lines, image.size)
  41
  42     if show_all:
  43         image_g = image.copy()
  44         draw = ImageDraw.Draw(image_g)
  45         for line in intersections:
  46             for (x, y) in line:
  47                 draw.point((x , y), fill=(120, 255, 120))
  48         do_something(image_g, "intersections")
  49
  50     return intersections
  51
  52 def board(image, intersections, show_all, do_something, logger):
  53     """Find stone colors and return board situation."""
  54
  55 #    image_c = filters.color_enhance(image)
  56 #    if show_all:
  57 #        do_something(image_c, "white balance")
  58     image_c = image
  59
  60     image_l = image_c.load()
  61     import Image, sys
  62     new_image = Image.new('RGB', (19 * 7, 19 * 7))
  63     image_nl = new_image.load()
  64     new_image2 = Image.new('L', (19 * 7, 19 * 7))
  65     image_nll = new_image2.load()
  66     y = 3
  67     for line in intersections:
  68         x = 3
  69         for (xi, yi) in line:
  70             for xx in [-3,-2,-1,0,1,2,3]:
  71                 for yy in [-3,-2,-1,0,1,2,3]:
  72                     try:
  73                         image_nl[x + xx, y + yy] = image_l[xi + xx, yi + yy]
  74                     except IndexError:
  75                         pass
  76             for xx in [-2,-1,0,1,2]:
  77                 for yy in [-2,-1,0,1,2]:
  78                     try:
  79                         z = xi + xx
  80                         w = yi + yy
  81                         luma = lambda ((r,g,b)): colorsys.rgb_to_hls(r / 255., g /
  82                                                                    255. ,b /
  83                                                                    255.)[1]
  84                         image_nll[x + xx, y + yy] = (luma(image_l[z, w]) * (-8)  +
  85                                                      luma(image_l[z - 1, w - 1]) +
  86                                                      luma(image_l[z - 1, w]) +
  87                                                      luma(image_l[z - 1, w + 1]) +
  88                                                      luma(image_l[z, w - 1]) +
  89                                                      luma(image_l[z, w + 1]) +
  90                                                      luma(image_l[z + 1, w - 1]) +
  91                                                      luma(image_l[z + 1, w]) +
  92                                                      luma(image_l[z + 1, w + 1])) * 255
  93                     except IndexError:
  94                         pass
  95             x += 7
  96         y += 7
  97     do_something(new_image, "intersections")
  98     do_something(new_image2, "intersections")
  99
 100
 101     board_raw = []
 102
 103     for line in intersections:
 104         board_raw.append([stone_color_raw(image_c, intersection) for intersection in
 105                       line])
 106     board_raw = sum(board_raw, [])
 107
 108     ### Show color distribution
 109
 110     if show_all:
 111         import matplotlib.pyplot as pyplot
 112         import Image
 113         fig = pyplot.figure(figsize=(8, 6))
 114         luma = [s[0] for s in board_raw]
 115         saturation = [s[1] for s in board_raw]
 116         pyplot.scatter(luma, saturation,
 117                        color=[s[2] for s in board_raw])
 118         pyplot.xlim(0,1)
 119         pyplot.ylim(0,1)
 120         fig.canvas.draw()
 121         size = fig.canvas.get_width_height()
 122         buff = fig.canvas.tostring_rgb()
 123         image_p = Image.fromstring('RGB', size, buff, 'raw')
 124         do_something(image_p, "color distribution")
 125
 126     #max_s0 = max(s[0] for s in board_raw)
 127     #min_s0 = min(s[0] for s in board_raw)
 128     #norm_s0 = lambda x: (x - min_s0) / (max_s0 - min_s0)
 129     #max_s1 = max(s[1] for s in board_raw)
 130     #min_s1 = min(s[1] for s in board_raw)
 131     #norm_s1 = lambda x: (x - min_s1) / (max_s1 - min_s1)
 132     #max_s1 = max(s[1] for s in board_raw)
 133     #min_s1 = min(s[1] for s in board_raw)
 134     #norm_s1 = lambda x: (x - min_s1) / (max_s1 - min_s1)
 135     #color_data = [(norm_s0(s[0]), norm_s1(s[1])) for s in board_raw]
 136     color_data = [(s[0], s[1],s[4]) for s in board_raw]
 137
 138     clusters, score = k_means.cluster(3, 3,zip(color_data, range(len(color_data))),
 139                                [[0., 0.5,0.0], [0.5, 0.5, 0.], [1., 0.5, 0.]])
 140 #    clusters1, score1 = k_means.cluster(1, 2,zip(color_data, range(len(color_data))),
 141 #                               [[0.5, 0.5]])
 142 #    clusters2, score2 = k_means.cluster(2, 2,zip(color_data, range(len(color_data))),
 143 #                               [[0., 0.5], [0.75, 0.5]])
 144 #    import sys
 145 #    print >> sys.stderr, score1, score2, score
 146 #
 147     if show_all:
 148         fig = pyplot.figure(figsize=(8, 6))
 149         pyplot.scatter([d[0][0] for d in clusters[0]], [d[0][1] for d in clusters[0]],
 150                                                  color=(1,0,0,1))
 151         pyplot.scatter([d[0][0] for d in clusters[1]], [d[0][1] for d in clusters[1]],
 152                                                  color=(0,1,0,1))
 153         pyplot.scatter([d[0][0] for d in clusters[2]], [d[0][1] for d in clusters[2]],
 154                                                  color=(0,0,1,1))
 155         pyplot.xlim(0,1)
 156         pyplot.ylim(0,1)
 157         fig.canvas.draw()
 158         size = fig.canvas.get_width_height()
 159         buff = fig.canvas.tostring_rgb()
 160         image_p = Image.fromstring('RGB', size, buff, 'raw')
 161         do_something(image_p, "color clustering")
 162
 163     clusters[0] = [(p[1], 'B') for p in clusters[0]]
 164     clusters[1] = [(p[1], '.') for p in clusters[1]]
 165     clusters[2] = [(p[1], 'W') for p in clusters[2]]
 166
 167     board_rl = sum(clusters, [])
 168     board_rl.sort()
 169     board_rg = (p[1] for p in board_rl)
 170
 171     board_r = []
 172
 173     #TODO 19 should be a size parameter
 174     try:
 175         for i in xrange(19):
 176             for _ in xrange(19):
 177                 board_r.append(board_rg.next())
 178     except StopIteration:
 179         pass
 180
 181     return output.Board(19, board_r)
 182
 183 def mean_luma(cluster):
 184     """Return mean luminanace of the *cluster* of points."""
 185     return sum(c[0][0] for c in cluster) / float(len(cluster))
 186
 187 def to_general(line, size):
 188     # TODO comment
 189     (x1, y1), (x2, y2) = linef.line_from_angl_dist(line, size)
 190     return (y2 - y1, x1 - x2, x2 * y1 - x1 * y2)
 191
 192 def intersection(l1, l2):
 193     a1, b1, c1 = l1
 194     a2, b2, c2 = l2
 195     delim = float(a1 * b2 - b1 * a2)
 196     x = (b1 * c2 - c1 * b2) / delim
 197     y = (c1 * a2 - a1 * c2) / delim
 198     return x, y
 199
 200 # TODO remove the parameter get_all
 201 def intersections_from_angl_dist(lines, size, get_all=True):
 202     """Take grid-lines and size of the image. Return intersections."""
 203     lines0 = map(lambda l: to_general(l, size), lines[0])
 204     lines1 = map(lambda l: to_general(l, size), lines[1])
 205     intersections = []
 206     for l1 in lines1:
 207         line = []
 208         for l2 in lines0:
 209             line.append(intersection(l1, l2))
 210         intersections.append(line)
 211     return intersections
 212
 213 def rgb2lumsat(color):
 214     """Convert RGB to luminance and HSI model saturation."""
 215     r, g, b = color
 216     luma = (0.30 * r + 0.59 * g + 0.11 * b) / 255.0
 217     max_diff = max(color) - min(color)
 218     if max_diff == 0:
 219         saturation = 0
 220     else:
 221         saturation = 1. - ((3. * min(color)) / sum(color))
 222     return luma, saturation
 223
 224 def median(lst):
 225     #TODO comment (or delete maybe?)
 226     len_lst = len(lst)
 227     if len_lst % 2 == 0:
 228         return (lst[len_lst / 2] + lst[len_lst / 2 + 1]) / 2.0
 229     else:
 230         return lst[len_lst / 2]
 231
 232 def stone_color_raw(image, (x, y)):
 233     """Given image and coordinates, return stone color."""
 234     size = 3
 235     points = []
 236     for i in range(-size, size + 1):
 237         for j in range(-size, size + 1):
 238             try:
 239                 points.append(image.getpixel((x + i, y + j)))
 240             except IndexError:
 241                 pass
 242     norm = float(len(points))
 243     if norm == 0:
 244         return 0, 0, (0, 0, 0) #TODO trow exception here
 245     norm = float(norm*255)
 246     color = (sum(p[0] for p in points) / norm,
 247              sum(p[1] for p in points) / norm,
 248              sum(p[2] for p in points) / norm)
 249     hue, luma, saturation = colorsys.rgb_to_hls(*color)
 250     color = colorsys.hls_to_rgb(hue, 0.5, 1.)
 251
 252     der = 0
 253     image_l = image.load()
 254     for xx in [-2,-1,0,1,2]:
 255         for yy in [-2,-1,0,1,2]:
 256             try:
 257                 z = x + xx
 258                 w = y + yy
 259                 lumal = lambda ((r,g,b)): colorsys.rgb_to_hls(r / 255., g /
 260                                                            255. ,b /
 261                                                            255.)[1]
 262                 der += (lumal(image_l[z, w]) * (-8)  +
 263                          lumal(image_l[z - 1, w - 1]) +
 264                          lumal(image_l[z - 1, w]) +
 265                          lumal(image_l[z - 1, w + 1]) +
 266                          lumal(image_l[z, w - 1]) +
 267                          lumal(image_l[z, w + 1]) +
 268                          lumal(image_l[z + 1, w - 1]) +
 269                          lumal(image_l[z + 1, w]) +
 270                          lumal(image_l[z + 1, w + 1]))
 271             except IndexError:
 272                 pass
 273
 274     der = max(der / 36., 0)
 275     return luma, saturation, color, hue, der