Implementando k-means clustering para a extração de cores

Sejam bem-vindos a mais um Turing Talks sobre Visão Computacional! Hoje apresentaremos a extração da paleta de cores de uma imagem a partir de k-means clustering, um algoritmo de clusterização muito simples e popular.

Você já deve ter visto alguma aplicação desse algoritmo mesmo sem saber. Ele é utilizado para gerar as paletas de cores que ficaram populares na internet recentemente.

Ao final desse texto, teremos apresentado o funcionamento do k-means, sua implementação e a extração da paleta de cores de uma imagem.

O que é clusterização

Algoritmos de clusterização visam a separar um conjunto de dados em grupos, conhecidos por clusters, de forma que os elementos de um mesmo agrupamento apresentem alguma característica semelhante. Cada um desses clusters possui um centroide, ou seja, um elemento central que representa seu respectivo grupo. No final, queremos maximizar essa diferença para que os grupos sejam o mais distintos possível.

K-means clustering

O método k-means clustering é um algoritmo de aprendizado de máquina não supervisionado bastante prático e simples de se entender. Primeiramente, a quantidade de clusters K deve ser determinada (daí que vem o ”k” em seu nome). Então, selecionamos K elementos dentro do nosso conjunto de dados aleatoriamente e estes serão os nossos centroides iniciais.

O próximo passo é realizar iterações até que encontremos os centroides corretos. Estas iterações seguem a seguinte ideia:

• Para cada dado, calculamos a distância euclidiana entre ele e as centroides e, assim, ele é atribuído ao cluster cuja distância entre o seu centroide e o elemento é menor.
• Uma vez que todos os elementos já estão agrupados em K grupos, recalculamos os seus centroides como a média dos pontos que os compõem (esta é a razão para “means”, ou “médias” em português, no nome do algoritmo).

Obtidos os novos centroides, podemos fazer mais uma iteração. Esse processo deve ocorrer indefinidamente até que a variação entre os centroides antigos e os novos seja muito pequena ou até que um número específico de iterações seja atingido.

Implementação

Nossa implementação vai ser construída em cima de imagens com três canais de cores, RGB (veja mais sobre aqui). Isso tem algumas implicações para o funcionamento do algoritmo.

• A fórmula para a distância Euclidiana D será a seguinte:

Sendo C referente ao centroide, E ao píxel analisado e R, G e B, as coordenadas de cores (canal vermelho, verde e azul).

• E interpretaremos os canais de cores de cada píxel como sendo coordenadas de um vetor com três dimensões.
• As fórmulas das coordenadas da centroide, obtidas a partir das médias das coordenadas dos N elementos que compõem o cluster são:

Agora, sabendo o que está por trás do algoritmo e da implementação, podemos começar a escrever nosso código em python.

Em primeiro lugar, vamos importar todas as bibliotecas que precisamos:

import cv2
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from mpl_toolkits.mplot3d import Axes3D
from PIL import Image
import pandas as pd
import colorsys

As principais funções criadas para esta implementação são apresentadas abaixo. Elas realizam todo o processo de iteração enunciado. O código completo se encontra em nosso Github.

def clusterData(centroids, data):
    """
    Realiza o processo de reagrupamento dos pontos: 
    - Obtem a distância entre cada ponto e os centroides
    - Agrupa os pontos com o centroide de menor distância
    """
    clusters = emptyMatrix(len(centroids))
    for i in data:
        distances = [distanceRGB(i, centroid) for centroid in centroids]
        clusters[np.argmin(distances)].append(i)
    return clusters
    
def computeCentroids(clusters):
    """Recalcula os centroides de cada agrupamento"""
    centroids = []
    for cluster in clusters:
        centroids.append(averageRGB(cluster))
    return np.array(centroids)
    
def kMeansClustering(k, data):
    """
    Realiza todo o processo de clusterização:
    - Gera centroides aleatórios
    - Iterações: 
        Atribuir pontos a cada centroide para criar agrupamentos
        Recalcular centroides a partir da média de cada agrupamento 
    """
    oldCentroids = randomCentroids(k, data)
    clusteredData = clusterData(oldCentroids, data)
    newCentroids = computeCentroids(clusteredData)
    
    count = 0
    
    while not isStable(oldCentroids, newCentroids) and count < 13:
        
        oldCentroids = newCentroids
        clusteredData = clusterData(newCentroids, data) 
        newCentroids = computeCentroids(clusteredData)

        count += 1
        
    return newCentroids, clusteredData, count

Visualização

Com o algoritmo implementado, podemos enfim vê-lo em ação em uma imagem. Para isso, precisamos importá-la e mudar suas dimensões para rodarmos o algoritmo em menos tempo.

Assim, estamos quase prontos para gerar nossa paleta. Mas antes, é interessante criarmos uma visualização que facilite nosso trabalho. Mencionamos anteriormente que as cores RGB da imagem podem ser interpretadas como coordenadas de um vetor. Se mapearmos o vetor equivalente a cada píxel no espaço, obtemos uma boa maneira de entender como o algoritmo vai funcionar.

s = 10
alpha = 1
lw = 0.05

fig = plt.figure(figsize = (8, 8))
ax = fig.add_subplot(111, projection = '3d')

ax.scatter(r, g, b, s = s, c = c, alpha = alpha,
           edgecolor = 'k', lw = lw)
ax.set(xlim3d = [0, 255], ylim3d = [0, 255], zlim3d = [0, 255])
ax.set_xlabel('Vermelho', fontsize = 14)
ax.set_ylabel('Verde', fontsize = 14)
ax.set_zlabel('Azul', fontsize = 14)
ax.set_title('Projeção de cores 3D', fontsize = 22)

plt.show()

A partir dessa projeção 3D, observamos a distribuição das cores da imagem. Há uma grande concentração de cores quentes, como amarelo, laranja e vermelho. Além disso, muitos pontos roxos e também azuis. Dessa forma, podemos escolher uma quantidade de clusters condizente. Portanto, utilizaremos K = 5.

Aplicação

Então, podemos aplicar o método em nossos dados, obtendo os centroides dos clusters, seus respectivos elementos e a quantidade de iterações necessárias até a convergência.

size = (img.shape[0]//10, img.shape[1])
img_plt = np.concatenate([img, criaPaleta(size, centros)])

fig = plt.figure(figsize = (20,20))
plt.imshow(img_plt)
plt.axis('off')
plt.show()

Percebemos que os centroides encontrados pelo algoritmo são próximos às cores que apontamos durante a visualização. Quanto às iterações, vemos os passos até a convergência com base na evolução dos centroides calculados no GIF a seguir:

Conclusão

Nesse Turing Talks vimos que o k-means é um algoritmo bastante eficiente para essa aplicação, apesar da sua simplicidade. Pelo GIF, notamos que são necessárias aproximadamente 8 iterações para chegar a um bom resultado.

Não se esqueça de seguir o Grupo Turing no Medium, Facebook, Instagram e LinkedIn. Além disso, faça parte do nosso servidor do Discord, Comunidade Turing dedicado à disseminação e ao estudo de Inteligência Artificial.

Autores: Rafael Araujo Coelho e Luís Henrique de Almeida Fernandes

Fonte:  https://link.medium.com/78CbwPTG3ab