save figures from matplotlib

[imago.git] / imago_pack / intrsc.py

diff --git a/imago_pack/intrsc.py b/imago_pack/intrsc.py
index d26bea4..1d4dd67 100644 (file)
--- a/imago_pack/intrsc.py
+++ b/imago_pack/intrsc.py
@@ -1,10 +1,11 @@
-"""Imago intersections module"""
+"""Imago intersections module."""

 
 from math import cos, tan, pi
 from operator import itemgetter
 
 import ImageDraw
 
 
 from math import cos, tan, pi
 from operator import itemgetter
 
 import ImageDraw
 

+import filters

 import k_means
 import output
 
 import k_means
 import output
 


@@ -25,6 +26,7 @@ def dst_sort(lines):
 
 def board(image, lines, show_all, do_something):
     """Compute intersections, find stone colors and return board situation."""
 
 def board(image, lines, show_all, do_something):
     """Compute intersections, find stone colors and return board situation."""

+    # TODO refactor show_all, do_something

     lines = [dst_sort(l) for l in lines]
     intersections = intersections_from_angl_dist(lines, image.size)
 
     lines = [dst_sort(l) for l in lines]
     intersections = intersections_from_angl_dist(lines, image.size)
 


@@ -36,10 +38,14 @@ def board(image, lines, show_all, do_something):
                 draw.point((x , y), fill=(120, 255, 120))
         do_something(image_g, "intersections")
 
                 draw.point((x , y), fill=(120, 255, 120))
         do_something(image_g, "intersections")
 

+    image_c = filters.color_enhance(image)
+    if show_all:
+        do_something(image_c, "white balance")
+    

     board_raw = []
     
     for line in intersections:
     board_raw = []
     
     for line in intersections:

-        board_raw.append([stone_color_raw(image, intersection) for intersection in
+        board_raw.append([stone_color_raw(image_c, intersection) for intersection in

                       line])
     board_raw = sum(board_raw, [])
 
                       line])
     board_raw = sum(board_raw, [])
 


@@ -49,23 +55,40 @@ def board(image, lines, show_all, do_something):
 
     if show_all:
         import matplotlib.pyplot as pyplot
 
     if show_all:
         import matplotlib.pyplot as pyplot

+        import Image
+        fig = pyplot.figure(figsize=(8, 6))

         pyplot.scatter(luma, saturation, 
         pyplot.scatter(luma, saturation, 

-                       color=[(s[2][0]/255., s[2][1]/255., s[2][2]/255., 1.) 
-                              for s in board_raw])
-        pyplot.show()
+                       color=[(s[2][0]/255.,
+                               s[2][1]/255.,
+                               s[2][2]/255., 1.) 
+                                   for s in board_raw])
+        pyplot.xlim(0,1)
+        pyplot.ylim(0,1)
+        fig.canvas.draw()
+        size = fig.canvas.get_width_height()
+        buff = fig.canvas.tostring_rgb()
+        image_p = Image.fromstring('RGB', size, buff, 'raw')
+        do_something(image_p, "color distribution")

 
     clusters = k_means.cluster(3, 2,zip(zip(luma, saturation), range(len(luma))),
 
     clusters = k_means.cluster(3, 2,zip(zip(luma, saturation), range(len(luma))),

-                               [[0., 0.], [0.5, 0.5], [1., 1.]])
+                               [[0., 0.5], [0.5, 0.5], [1., 0.5]])

    #clusters.sort(key=mean_luma)
 
     if show_all:
    #clusters.sort(key=mean_luma)
 
     if show_all:

+        fig = pyplot.figure(figsize=(8, 6))

         pyplot.scatter([d[0][0] for d in clusters[0]], [d[0][1] for d in clusters[0]],
                                                  color=(1,0,0,1))
         pyplot.scatter([d[0][0] for d in clusters[1]], [d[0][1] for d in clusters[1]],
                                                  color=(0,1,0,1))
         pyplot.scatter([d[0][0] for d in clusters[2]], [d[0][1] for d in clusters[2]],
                                                  color=(0,0,1,1))
         pyplot.scatter([d[0][0] for d in clusters[0]], [d[0][1] for d in clusters[0]],
                                                  color=(1,0,0,1))
         pyplot.scatter([d[0][0] for d in clusters[1]], [d[0][1] for d in clusters[1]],
                                                  color=(0,1,0,1))
         pyplot.scatter([d[0][0] for d in clusters[2]], [d[0][1] for d in clusters[2]],
                                                  color=(0,0,1,1))

-        pyplot.show()
+        pyplot.xlim(0,1)
+        pyplot.ylim(0,1)
+        fig.canvas.draw()
+        size = fig.canvas.get_width_height()
+        buff = fig.canvas.tostring_rgb()
+        image_p = Image.fromstring('RGB', size, buff, 'raw')
+        do_something(image_p, "color clustering")

 
     clusters[0] = [(p[1], 'B') for p in clusters[0]]
     clusters[1] = [(p[1], '.') for p in clusters[1]]
 
     clusters[0] = [(p[1], 'B') for p in clusters[0]]
     clusters[1] = [(p[1], '.') for p in clusters[1]]


@@ -89,6 +112,7 @@ def board(image, lines, show_all, do_something):
     return output.Board(19, board_r)
 
 def mean_luma(cluster):
     return output.Board(19, board_r)
 
 def mean_luma(cluster):

+    """Return mean luma of the *cluster* of points."""

     return sum(c[0][0] for c in cluster) / float(len(cluster))
 
 def intersections_from_angl_dist(lines, size, get_all=True):
     return sum(c[0][0] for c in cluster) / float(len(cluster))
 
 def intersections_from_angl_dist(lines, size, get_all=True):


@@ -108,19 +132,29 @@ def intersections_from_angl_dist(lines, size, get_all=True):
         intersections.append(line)
     return intersections
    
         intersections.append(line)
     return intersections
    

+def RGBtoSat(c):
+    """Using the HSI color model."""
+    max_diff = max(c) - min(c)
+    if max_diff == 0:
+        return 0
+    else:
+        return 1. - ((3. * min(c)) / sum(c)) 
+

 def stone_color_raw(image, (x, y)):
     """Given image and coordinates, return stone color."""
 def stone_color_raw(image, (x, y)):
     """Given image and coordinates, return stone color."""

+    size = 3 

     suma = []
     suma = []

-    for i in range(-2, 3):
-        for j in range(-2, 3):
+    t = 0
+    for i in range(-size, size + 1):
+        for j in range(-size, size + 1):

             try:
                 suma.append(image.getpixel((x + i, y + j)))
             try:
                 suma.append(image.getpixel((x + i, y + j)))

+                t += 1

             except IndexError:
                 pass
             except IndexError:
                 pass

-    luma = sum([0.30 * sum(s[0] for s in suma) / 25., 0.59 * sum(s[1] for s in suma) / 25., 
-            0.11 * sum(s[2] for s in suma) / 25.]) / 255.
-    saturation = sum(max(s) - min(s) / float(255. - abs(max(s) + min(s) - 255))
-                     for s in suma) / (25. * 255.)
-    color = [sum(s[0] for s in suma) / 25., sum(s[1] for s in suma) / 25.,
-             sum(s[2] for s in suma) / 25.]
+    luma = sum([0.30 * sum(s[0] for s in suma) / t, 0.59 * sum(s[1] for s in suma) / t, 
+            0.11 * sum(s[2] for s in suma) / t]) / 255.
+    saturation = sum(RGBtoSat(s) for s in suma) / t
+    color = [sum(s[0] for s in suma) / t, sum(s[1] for s in suma) / t,
+             sum(s[2] for s in suma) / t]

     return luma, saturation, color
     return luma, saturation, color