imago_pack/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 k_means
   9
  10 def dst(line):
  11     """Return normalized line."""
  12     if line[0] < pi / 2:
  13         line = line[0] + pi, - line[1]
  14     return line
  15
  16 def dst_sort(lines):
  17     """Return lines sorted by distance."""
  18     l_max = max(l[0] for l in lines)
  19     l_min = min(l[0] for l in lines)
  20     if l_max - l_min > (3. / 4) * pi:
  21         lines = [dst(l) for l in lines]
  22     lines.sort(key=itemgetter(1))
  23     return lines
  24
  25 def board(image, lines, show_all, do_something):
  26     """Compute intersections, find stone colors and return board situation."""
  27     lines = [dst_sort(l) for l in lines]
  28     intersections = intersections_from_angl_dist(lines, image.size)
  29
  30     if show_all:
  31         image_g = image.copy()
  32         draw = ImageDraw.Draw(image_g)
  33         for line in intersections:
  34             for (x, y) in line:
  35                 draw.point((x , y), fill=(120, 255, 120))
  36         do_something(image_g, "intersections")
  37
  38     board_raw = []
  39
  40     for line in intersections:
  41         board_raw.append([stone_color_raw(image, intersection) for intersection in
  42                       line])
  43     board_raw = sum(board_raw, [])
  44
  45     ### Show color distribution
  46     luma = [s[0] for s in board_raw]
  47     saturation = [s[1] for s in board_raw]
  48
  49     if show_all:
  50         import matplotlib.pyplot as pyplot
  51         pyplot.scatter(luma, saturation,
  52                        color=[(s[2][0]/255., s[2][1]/255., s[2][2]/255., 1.)
  53                               for s in board_raw])
  54         pyplot.show()
  55
  56     clusters = k_means.cluster(3, 2,zip(zip(luma, saturation), range(len(luma))),
  57                                [[0., 0.], [0.5, 0.5], [1., 1.]])
  58    #clusters.sort(key=mean_luma)
  59
  60     if show_all:
  61         pyplot.scatter([d[0][0] for d in clusters[0]], [d[0][1] for d in clusters[0]],
  62                                                  color=(1,0,0,1))
  63         pyplot.scatter([d[0][0] for d in clusters[1]], [d[0][1] for d in clusters[1]],
  64                                                  color=(0,1,0,1))
  65         pyplot.scatter([d[0][0] for d in clusters[2]], [d[0][1] for d in clusters[2]],
  66                                                  color=(0,0,1,1))
  67         pyplot.show()
  68
  69     clusters[0] = [(p[1], 'B') for p in clusters[0]]
  70     clusters[1] = [(p[1], '.') for p in clusters[1]]
  71     clusters[2] = [(p[1], 'W') for p in clusters[2]]
  72
  73     board_rl = sum(clusters, [])
  74     board_rl.sort()
  75     board_rg = (p[1] for p in board_rl)
  76
  77     board_r = []
  78
  79     try:
  80         for i in xrange(19):
  81             board_r.append([])
  82             for _ in xrange(19):
  83                 board_r[i].append(board_rg.next())
  84     except StopIteration:
  85         pass
  86
  87     return board_r
  88
  89 def mean_luma(cluster):
  90     return sum(c[0][0] for c in cluster) / float(len(cluster))
  91
  92 def intersections_from_angl_dist(lines, size, get_all=True):
  93     """Take grid-lines and size of the image. Return intersections."""
  94     intersections = []
  95     for (angl1, dist1) in lines[1]:
  96         line = []
  97         for (angl2, dist2) in lines[0]:
  98             if abs(angl1 - angl2) > 0.4:
  99                 i_x =  (- ((dist2 / cos(angl2)) - (dist1 / cos(angl1)))
 100                         / (tan(angl1) - tan(angl2)))
 101                 i_y = (tan(angl1) * i_x) - (dist1 / cos(angl1))
 102                 if get_all or (-size[0] / 2 < i_x < size[0] / 2 and
 103                     -size[1] / 2 < i_y < size[1] / 2):
 104                     line.append((int(i_x + size[0] / 2),
 105                                  int(i_y + size[1] / 2)))
 106         intersections.append(line)
 107     return intersections
 108
 109 def stone_color_raw(image, (x, y)):
 110     """Given image and coordinates, return stone color."""
 111     suma = []
 112     for i in range(-2, 3):
 113         for j in range(-2, 3):
 114             try:
 115                 suma.append(image.getpixel((x + i, y + j)))
 116             except IndexError:
 117                 pass
 118     luma = sum([0.30 * sum(s[0] for s in suma) / 25., 0.59 * sum(s[1] for s in suma) / 25.,
 119             0.11 * sum(s[2] for s in suma) / 25.]) / 255.
 120     saturation = sum(max(s) - min(s) / float(255. - abs(max(s) + min(s) - 255))
 121                      for s in suma) / (25. * 255.)
 122     color = [sum(s[0] for s in suma) / 25., sum(s[1] for s in suma) / 25.,
 123              sum(s[2] for s in suma) / 25.]
 124     return luma, saturation, color