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