-
Notifications
You must be signed in to change notification settings - Fork 0
/
colours.py
100 lines (71 loc) · 2.37 KB
/
colours.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
"""
K-means dominant colours
Usage: colours.py [-h] [-n N] filename
-n N Number of dominant colours [default: 3].
"""
from collections import namedtuple
from math import sqrt
import random
import argparse
from PIL import Image
Point = namedtuple("Point", ("coords", "n", "ct"))
Cluster = namedtuple("Cluster", ("points", "center", "n"))
def get_points(img):
points = []
w, h = img.size
for count, color in img.getcolors(w * h):
points.append(Point(color[:3], 3, count))
return points
def rgb_to_hex(rgb):
return "#%02x%02x%02x" % tuple(rgb)
def colorz(filename, n=3):
img = Image.open(filename)
img.thumbnail((200, 200))
points = get_points(img)
clusters = kmeans(points, n, 1)
rgbs = [map(int, c.center.coords) for c in clusters]
return map(rgb_to_hex, rgbs)
def euclidean(p1, p2):
return sqrt(sum([(p1.coords[i] - p2.coords[i]) ** 2 for i in range(p1.n)]))
def calculate_center(points, n):
vals = [0.0] * n
plen = 0
for p in points:
plen += p.ct
for i in range(n):
vals[i] += p.coords[i] * p.ct
return Point([(v / plen) for v in vals], n, 1)
def kmeans(points, k, min_diff):
clusters = [Cluster([p], p, p.n) for p in random.sample(points, k)]
while 1:
plists = [[] for _ in range(k)]
for p in points:
idx = 0
smallest_distance = float("Inf")
for i in range(k):
distance = euclidean(p, clusters[i].center)
if distance < smallest_distance:
smallest_distance = distance
idx = i
plists[idx].append(p)
diff = 0
for i in range(k):
old = clusters[i]
center = calculate_center(plists[i], old.n)
new = Cluster(plists[i], center, old.n)
clusters[i] = new
diff = max(diff, euclidean(old.center, new.center))
if diff < min_diff:
break
return clusters
def main():
parser = argparse.ArgumentParser(description="K-means dominant colours")
parser.add_argument("filename", help="The filename of the image")
parser.add_argument(
"-n", type=int, default=3, help="Number of dominant colours (default: 3)"
)
args = parser.parse_args()
for colour in colorz(args.filename, args.n):
print(colour)
if __name__ == "__main__":
main()