Определение доминирующих цветов: Python и метод k-средних

^©Assorium

На Хабре публиковалось несколько статей с алгоритмами и скриптами для выбора доминирующих цветов на изображении: 1, 2, 3. В комментариях к тем статьям можно найти ссылки ещё на десяток подобных программ и сервисов. Но нет предела совершенству — и почему бы не рассмотреть способ, который кажется самым оптимальным? Речь идёт об использовании кластеризации методом k-средних (k-means).

Как и многие до него, американский веб-разработчик Чарльз Лейфер (Charles Leifer) использовал метод k-средних для кластеризации цветов на изображении. Идея метода при кластеризации любых данных заключается в том, чтобы минимизировать суммарное квадратичное отклонение точек кластеров от центров этих кластеров. На первом этапе выбираются случайным образом начальные точки (центры масс) и вычисляется принадлежность каждого элемента к тому или иному центру. Затем на каждой итерации выполнения алгоритма происходит перевычисление центров масс — до тех пор, пока алгоритм не сходится.

В результате получается примерно такая картина. Точки раскрашены, в зависимости от цвета кластера, чёрные точки отображают центры масс.

В применении к изображениям каждый пиксель позиционируется в трёхмерном пространстве RGB, где вычисляется расстояние до центров масс. Для оптимизации картинки уменьшаются до 200х200 с помощью библиотеки PIL. Она же используется для извлечения значений RGB.

Код

from collections import namedtuple
from math import sqrt
import random
try:
    import Image
except ImportError:
    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, count))
    return points

rtoh = lambda rgb: '#%s' % ''.join(('%02x' % p for p in rgb))

def colorz(filename, n=3):
    img = Image.open(filename)
    img.thumbnail((200, 200))
    w, h = img.size

    points = get_points(img)
    clusters = kmeans(points, n, 1)
    rgbs = [map(int, c.center.coords) for c in clusters]
    return map(rtoh, 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 for i in range(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 i in range(k)]

        for p in points:
            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

Примеры

Определение доминирующих цветов — довольно полезная вещь, которой всегда найдётся применение. Это выбор палитры для веб-сайта или некоторых элементов UI. Например, браузер Chrome использует метод k-средних для выбора доминирующего цвета с фавикона.

Автор: alizar