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 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     lines = [dst_sort(l) for l in lines]
  31     intersections = intersections_from_angl_dist(lines, image.size)
  32
  33     if show_all:
  34         image_g = image.copy()
  35         draw = ImageDraw.Draw(image_g)
  36         for line in intersections:
  37             for (x, y) in line:
  38                 draw.point((x , y), fill=(120, 255, 120))
  39         do_something(image_g, "intersections")
  40
  41     image_c = filters.color_enhance(image)
  42     if show_all:
  43         do_something(image_c, "white balance")
  44
  45     board_raw = []
  46
  47     for line in intersections:
  48         board_raw.append([stone_color_raw(image_c, intersection) for intersection in
  49                       line])
  50     board_raw = sum(board_raw, [])
  51
  52     ### Show color distribution
  53     luma = [s[0] for s in board_raw]
  54     saturation = [s[1] for s in board_raw]
  55
  56     if show_all:
  57         import matplotlib.pyplot as pyplot
  58         pyplot.scatter(luma, saturation,
  59                        color=[(s[2][0]/255.,
  60                                s[2][1]/255.,
  61                                s[2][2]/255., 1.)
  62                                    for s in board_raw])
  63         pyplot.xlim(0,1)
  64         pyplot.ylim(0,1)
  65         pyplot.show()
  66
  67     clusters = k_means.cluster(3, 2,zip(zip(luma, saturation), range(len(luma))),
  68                                [[0., 0.5], [0.5, 0.5], [1., 0.5]])
  69    #clusters.sort(key=mean_luma)
  70
  71     if show_all:
  72         pyplot.scatter([d[0][0] for d in clusters[0]], [d[0][1] for d in clusters[0]],
  73                                                  color=(1,0,0,1))
  74         pyplot.scatter([d[0][0] for d in clusters[1]], [d[0][1] for d in clusters[1]],
  75                                                  color=(0,1,0,1))
  76         pyplot.scatter([d[0][0] for d in clusters[2]], [d[0][1] for d in clusters[2]],
  77                                                  color=(0,0,1,1))
  78         pyplot.xlim(0,1)
  79         pyplot.ylim(0,1)
  80         pyplot.show()
  81
  82     clusters[0] = [(p[1], 'B') for p in clusters[0]]
  83     clusters[1] = [(p[1], '.') for p in clusters[1]]
  84     clusters[2] = [(p[1], 'W') for p in clusters[2]]
  85
  86     board_rl = sum(clusters, [])
  87     board_rl.sort()
  88     board_rg = (p[1] for p in board_rl)
  89
  90     board_r = []
  91
  92     #TODO 19 should be a size parameter
  93     try:
  94         for i in xrange(19):
  95             for _ in xrange(19):
  96                 board_r.append(board_rg.next())
  97     except StopIteration:
  98         pass
  99
 100
 101     return output.Board(19, board_r)
 102
 103 def mean_luma(cluster):
 104     """Return mean luma of the *cluster* of points."""
 105     return sum(c[0][0] for c in cluster) / float(len(cluster))
 106
 107 def intersections_from_angl_dist(lines, size, get_all=True):
 108     """Take grid-lines and size of the image. Return intersections."""
 109     intersections = []
 110     for (angl1, dist1) in lines[1]:
 111         line = []
 112         for (angl2, dist2) in lines[0]:
 113             if abs(angl1 - angl2) > 0.4:
 114                 i_x =  (- ((dist2 / cos(angl2)) - (dist1 / cos(angl1)))
 115                         / (tan(angl1) - tan(angl2)))
 116                 i_y = (tan(angl1) * i_x) - (dist1 / cos(angl1))
 117                 if get_all or (-size[0] / 2 < i_x < size[0] / 2 and
 118                     -size[1] / 2 < i_y < size[1] / 2):
 119                     line.append((int(i_x + size[0] / 2),
 120                                  int(i_y + size[1] / 2)))
 121         intersections.append(line)
 122     return intersections
 123
 124 def RGBtoSat(c):
 125     """Using the HSI color model."""
 126     max_diff = max(c) - min(c)
 127     if max_diff == 0:
 128         return 0
 129     else:
 130         return 1. - ((3. * min(c)) / sum(c))
 131
 132 def stone_color_raw(image, (x, y)):
 133     """Given image and coordinates, return stone color."""
 134     size = 3
 135     suma = []
 136     t = 0
 137     for i in range(-size, size + 1):
 138         for j in range(-size, size + 1):
 139             try:
 140                 suma.append(image.getpixel((x + i, y + j)))
 141                 t += 1
 142             except IndexError:
 143                 pass
 144     luma = sum([0.30 * sum(s[0] for s in suma) / t, 0.59 * sum(s[1] for s in suma) / t,
 145             0.11 * sum(s[2] for s in suma) / t]) / 255.
 146     saturation = sum(RGBtoSat(s) for s in suma) / t
 147     color = [sum(s[0] for s in suma) / t, sum(s[1] for s in suma) / t,
 148              sum(s[2] for s in suma) / t]
 149     return luma, saturation, color