X-Git-Url: http://git.tomasm.cz/imago.git/blobdiff_plain/e081eb1ecc06066d8015da7ad40947916607e964..97440043d1e48dfa0b1b197706866f5975cdb701:/imago.py

diff --git a/imago.py b/imago.py
index 1cda8ba..ff49826 100755
--- a/imago.py
+++ b/imago.py
@@ -47,6 +47,9 @@ def main():
     except IOError, msg:
         print >> sys.stderr, msg
         return 1
+    if image.mode == 'P':
+        image = image.convert('RGB')
+    
     if image.size[0] > args.w:
         image = image.resize((args.w, int((float(args.w)/image.size[0]) *
                               image.size[1])), Image.ANTIALIAS)
@@ -66,7 +69,7 @@ def main():
     if verbose:
         print >> sys.stderr, "edge detection"
 
-    im_edges = filters.edge_detection(im_l)
+    im_edges = filters.edge_detection(im_l, 2)
     if show_all:    
         do_something(im_edges, "edge detection")
 
@@ -90,14 +93,17 @@ def main():
     if show_all:
         do_something(im_h2, "second high pass filters")
 
-    im_h2 = filters.components(im_h2)
+    im_h2 = filters.components2(im_h2)
     if show_all:
         do_something(im_h2, "components centers")
 
     if verbose:
         print >> sys.stderr, "second hough transform"
 
-    hough2 = Hough(im_h2.size)
+    hough2 = Hough(im_h2.size) 
+    # im_hough might be used instead im_h2, but at the moment it brings a lot of
+    # noise to the second transform, which later confuses the center-finding
+    # mechanism (which is not very robust yet)
     im_hough2 = hough2.transform(im_h2)
     if show_all:
         do_something(im_hough2, "second hough transform")
@@ -106,26 +112,40 @@ def main():
     if show_all:
         do_something(im_h3, "third high pass filter")
      
-    im_h3 = filters.half_centers(im_h3)
+    im_h3 = filters.components(im_h3)
     if show_all:
         do_something(im_h3, "half centers")
 
     if verbose:
         print >> sys.stderr, "finding the grid"
 
-    lines_m = hough2.all_lines(im_h3)
+    lines_m = hough2.all_lines_h(im_h3)
     lines = []
+    im_c = im_h2.convert('RGB').convert('RGB', (1, 0.5, 0.5, 0))
+    draw_c = ImageDraw.Draw(im_c)
 
-    for line in lines_m:
+    for line_l in lines_m:
         im_line = Image.new('L', im_h2.size)
         draw = ImageDraw.Draw(im_line)
-        draw.line(line_from_angl_dist(line, im_h2.size), fill=255, width=5)
-        if show_all:
-            do_something(im_line, "line")
-        im_c = combine(im_h2, im_line)
-        if show_all:
-            do_something(im_c, "hough x lines")
-        lines.append(hough1.all_lines(im_c))
+        line_points = set()
+        for line in line_l:
+            draw.line(line_from_angl_dist(line, im_h2.size), fill=255, width=7)
+            draw_c.line(line_from_angl_dist(line, im_c.size), fill=(70, 70, 70), width=7)
+            for p in combine(im_h2, im_line):
+                line_points.add(p)
+        for point in line_points:
+            draw_c.point(point, fill=(120, 255, 120))
+        lines.append(hough1.lines_from_list(line_points))
+
+    if show_all:
+        do_something(im_c, "hough x lines")
+
+    image_g = image.copy()
+    draw = ImageDraw.Draw(image_g)
+    for line in [l for s in lines for l in s]:
+        draw.line(line_from_angl_dist(line, image.size), fill=(120, 255, 120))
+    if show_all:
+        do_something(image_g, "the grid")
 
     intersections = intersections_from_angl_dist(lines, image.size)
     image_g = image.copy()
@@ -139,7 +159,7 @@ def main():
                        line])
 
     if show_all:
-        do_something(image_g, "the grid")
+        do_something(image_g, "intersections")
 
     return 0
 
@@ -147,7 +167,10 @@ def stone_color(image, (x, y)):
     suma = 0.
     for i in range(-2, 3):
         for j in range(-2, 3):
-            suma += sum(image.getpixel((x + i, y + j)))
+            try:
+                suma += sum(image.getpixel((x + i, y + j)))
+            except IndexError:
+                pass
     suma /= 3 * 25
     if suma < 55:
         return 'B'
@@ -169,14 +192,13 @@ def combine(image1, image2):
     im_l1 = image1.load()
     im_l2 = image2.load()
 
-    im_n = Image.new('L', image1.size)
-    im_nl = im_n.load()
+    on_both = []
 
     for x in xrange(image1.size[0]):
         for y in xrange(image1.size[1]):
             if im_l1[x, y] and im_l2[x, y]:
-                im_nl[x, y] = 255
-    return im_n
+                on_both.append((x, y))
+    return on_both
 
 def line_from_angl_dist((angle, distance), size):
     x1 = - size[0] / 2
@@ -189,11 +211,13 @@ def intersections_from_angl_dist(lines, size):
     intersections = []
     for (angl1, dist1) in sorted(lines[1], key=itemgetter(1)):
         line = []
-        for (angl2, dist2) in lines[0]:
+        for (angl2, dist2) in sorted(lines[0], key=itemgetter(1)):
             if abs(angl1 - angl2) > 0.4:
                 x =  - ((dist2 / math.cos(angl2))-(dist1 / math.cos(angl1))) / (math.tan(angl1) - math.tan(angl2))
                 y = (math.tan(angl1) * x) - (dist1 / math.cos(angl1))
-                line.append((int(x + size[0] / 2), int(y + size[1] / 2)))
+                if (-size[0] / 2 < x < size[0] / 2 and 
+                    -size[1] / 2 < y < size[1] / 2):
+                    line.append((int(x + size[0] / 2), int(y + size[1] / 2)))
         intersections.append(line)
     return intersections