even more robust gridf3
[imago.git] / src / filters.py
1 """Image filters module.
2
3 All filters return a filtered copy of the image, the original image is
4 preserved.
5 """
6
7 from PIL import Image, ImageFilter
8
9 import pcf
10
11 def color_enhance(image):
12     """Stretch all color channels to their full range."""
13     image_l = image.load()
14     min_r, min_g, min_b = 999, 999, 999
15     max_r, max_g, max_b = -1, -1, -1
16
17     for x in xrange(image.size[0]):
18         for y in xrange(image.size[1]):
19             min_r = min(min_r, image_l[x, y][0])
20             max_r = max(max_r, image_l[x, y][0])
21             min_g = min(min_g, image_l[x, y][1])
22             max_g = max(max_g, image_l[x, y][1])
23             min_b = min(min_b, image_l[x, y][2])
24             max_b = max(max_b, image_l[x, y][2])
25
26     new_image = Image.new('RGB', image.size)
27     new_image_l = new_image.load()
28     for x in xrange(image.size[0]):
29         for y in xrange(image.size[1]):
30             r, g, b = image_l[x, y]
31             r = (r - min_r) * 255 / (max_r - min_r)
32             g = (g - min_g) * 255 / (max_g - min_g)
33             b = (b - min_b) * 255 / (max_b - min_b)
34             new_image_l[x, y] = (r, g, b)
35
36     return new_image
37
38 def edge_detection(image):
39     """Edge detection (on BW images)."""
40     new_image = image.filter(ImageFilter.GaussianBlur())
41     # GaussianBlur is undocumented class, it might not work in future versions
42     # of PIL
43     new_image = Image.fromstring('L', image.size,
44                              pcf.edge(image.size, image.tostring()))
45     return new_image
46
47 def peaks(image):
48     """Peak filter (on BW images)."""
49     image_l = image.load()
50     new_image = Image.new('L', image.size)
51     new_image_l = new_image.load()
52     
53     for x in range(2, image.size[0] - 2):
54         for y in range(2, image.size[1] - 2):
55             pix = (sum([sum([
56                 - image_l[a, b] 
57                     for b in range(y - 2, y + 3)]) 
58                     for a in range(x - 2, x + 3)])
59                 + (17 * image_l[x, y]))
60             if pix > 255:
61                 pix = 255
62             if pix < 0:
63                 pix = 0 
64             new_image_l[x, y] = pix
65     return new_image
66
67 def high_pass(image, height):
68     """High pass filter (on BW images)."""
69     image_l = image.load()
70     new_image = Image.new('L', image.size)
71     new_image_l = new_image.load()
72     
73     for x in xrange(image.size[0]):
74         for y in xrange(image.size[1]):
75             if image_l[x, y] < height:
76                 new_image_l[x, y] = 0
77             else:
78                 new_image_l[x, y] = image_l[x, y]
79
80     return new_image
81
82 def components(image, diameter):
83     # TODO comment 
84     # TODO refactor
85     image_l = image.load()
86     new_image = Image.new('L', image.size)
87     new_image_l = new_image.load()
88
89     components = [None]
90     comp_counter = 1
91
92     if diameter == 1:
93         for y in xrange(1, image.size[1] - 1):
94             for x in xrange(1, image.size[0] - 1):
95                 if image_l[x, y]:
96                     s = {0}
97                     s.add(new_image_l[x - 1, y - 1])
98                     s.add(new_image_l[x, y - 1])
99                     s.add(new_image_l[x + 1, y - 1])
100                     s.add(new_image_l[x - 1, y])
101                     if len(s) == 1:
102                         components.append(set())
103                         new_image_l[x, y] = comp_counter
104                         components[comp_counter].add((x, y))
105                         comp_counter += 1
106                     elif len(s) == 2:
107                         s.remove(0)
108                         c = s.pop()
109                         new_image_l[x, y] = c
110                         components[c].add((x,y))
111                     else:
112                         s.remove(0)
113                         c1, c2 = s.pop(), s.pop()
114                         components[c2].add((x, y))
115                         for (x1, y1) in components[c2]:
116                             new_image_l[x1, y1] = c1
117                         components[c1] = components[c1] | components[c2]
118                         components[c2] = None
119     elif diameter == 2:
120         for y in xrange(2, image.size[1] - 2):
121             for x in xrange(2, image.size[0] - 2):
122                 if image_l[x, y]:
123
124                     s = {0}
125                     for (a, b) in [(a,b) for a in range(x - 2, x + 3) 
126                               for b in range(y - 2, y + 1)]:
127                         if not (b == y and a >= x):
128                             s.add(new_image_l[a, b])
129
130                     if len(s) == 1:
131                         components.append(set())
132                         new_image_l[x, y] = comp_counter
133                         components[comp_counter].add((x, y))
134                         comp_counter += 1
135                     elif len(s) == 2:
136                         s.remove(0)
137                         c = s.pop()
138                         new_image_l[x, y] = c
139                         components[c].add((x,y))
140                     else:
141                         s.remove(0)
142                         c1 = s.pop()
143                         components[c1].add((x, y))
144                         new_image_l[x, y] = c1
145                         for c2 in s:
146                             for (x1, y1) in components[c2]:
147                                 new_image_l[x1, y1] = c1
148                             components[c1] = components[c1] | components[c2]
149                             components[c2] = None
150     else:
151         pass #TODO error
152
153
154     new_image = Image.new('L', image.size)
155     new_image_l = new_image.load()
156
157     for component in components:
158         if component:
159             x_c = 0
160             y_c = 0
161             c = 0
162             for (x, y) in component:
163                 x_c += x
164                 y_c += y
165                 c += 1
166             new_image_l[int(round(float(x_c)/c)), int(round(float(y_c)/c))] = 255
167
168     return new_image
169