Clustering K-Means para cuantización de imágenes

En este post exploramos la cuantización de imágenes aplicando el algoritmo k-means a una imagen de muestra de una flor. Vamos a ver cómo distintos valores de k afectan la calidad de la imagen resultante y la complejidad computacional del algoritmo. ¡Empecemos!

Módulos implementados

Todos los módulos de esta sección fueron implementados por el autor usando las librerías NumPy y Matplotlib. Puedes encontrar el código fuente en GitHub.

Documentación

src.image

• Image: representa una imagen con atributos como la matriz RGB, el vector RGB, la forma y el título.

• load_image(image_path: str) -> Image: carga una imagen desde la ruta indicada, la convierte a formato RGB y devuelve un objeto Image.

• rgb_vector_to_image(rgb_vector: np.ndarray, image_shape: tuple) -> Image: convierte un vector RGB en un objeto Image con la forma especificada.

• display_images(*images: Image, figsize=(10, 5)): muestra imágenes usando matplotlib.

• plot_rgb_vector(rgb_vector: np.ndarray, figsize=(10, 5), title=None, alpha=0.6, s=4): grafica un scatter plot 3D de vectores RGB.

from src import image

src.model

• KMeans: implementa el algoritmo de clustering \(k\)-means con parámetros para la cantidad de clusters (n_clusters), la tolerancia (tol), el máximo de iteraciones (max_iter) y el random state (random_state).

• KMeans.fit(X: np.ndarray) -> KMeans: ajusta el modelo KMeans a los datos de entrada X usando los parámetros especificados.

• KMeans.labels_: un array de forma (n_samples,) donde cada elemento representa el índice del cluster asignado a la muestra correspondiente.

• KMeans.cluster_centers_: un array de forma (n_clusters, n_features) que representa los centros finales de los clusters después de ajustar el modelo. Cada fila corresponde al centroide de un cluster en el espacio de features.

• plot_k_metrics(k_values: list[int], clusters_cost_list: list[float], execution_time_list: list[float], figsize: tuple = (10, 4)): grafica el costo de los clusters y el tiempo de ejecución de \(k\)-means para distintos valores de \(k\).

from src import model

Pruébalo tú mismo

Para probar estos módulos, sigue estos pasos:

1. Descarga el código fuente acá

2. Ve a la carpeta /image_quantization

3. Instala los paquetes necesarios: pip install -r requirements.txt

4. Crea un notebook o un archivo Python en la carpeta /image_quantization

5. ¡Listo! Ya puedes importar los módulos con from src import image, model

Cuantización de imágenes

Una de las aplicaciones de \(k\)-means es la cuantización de imágenes (más precisamente, la cuantización de color), que consiste en reducir la cantidad de colores distintos usados en una imagen, en general con la intención de que la imagen nueva sea lo más parecida posible a la original.

Exploremos este concepto aplicando el algoritmo \(k\)-means a una imagen de muestra (imagen de margarita, por Bessi).

Cargando la imagen

import numpy as np

img = image.load_image("img/flower.jpg")
img.title = "Original Image"
image.display_images(img)

La imagen se carga como una matriz \(M \in \mathbb{R}^{n \times m \times 3}\) donde \(M_{ij}\) es el píxel de la fila \(i\) y la columna \(j\):

img.rgb_matrix.shape

(417, 640, 3)

Para ajustar el algoritmo \(k\)-means, la matriz de la imagen \(M\) tiene que convertirse en una matriz de features \(X \in \mathbb{R}^{n \cdot m \times 3}\) donde \(X_i\) representa el vector de features \(x^{(i)}\):

img.rgb_vector.shape

(266880, 3)

Podemos visualizar los vectores de features en el espacio de 3 dimensiones:

image.plot_rgb_vector(img.rgb_vector, title="Feature vectors $$ X $$")

Vamos a experimentar con tres valores distintos de \(k\): 8, 16 y 32. Es decir que vamos a intentar representar todo el espectro de colores de la imagen usando 8, 16 o 32 colores distintos. El algoritmo \(k\)-means va a agrupar los colores parecidos y asignarlos al mismo cluster, reduciendo así la variedad de colores presentes en la imagen.

\(8\)-means

Primero entrenamos el modelo con \(8\) clusters.

k8_means = model.KMeans(n_clusters=8, random_state=44)
%time k8_means.fit(img.rgb_vector)

CPU times: total: 10.9 s
Wall time: 11.1 s

<src.model.KMeans at 0x1237c937210>

Podemos inspeccionar las iteraciones ejecutadas:

k8_means.iterations

72

Y podemos acceder a las etiquetas y los centroides (o centros) encontrados:

k8_means.labels_

array([6, 6, 6, …, 6, 6, 6], dtype=int64)

k8_means.cluster_centers_

array([
[42.39915168, 67.65893264, 84.11651684],
[ 87.46747399, 127.68052914, 119.33007048],
[191.76564391, 194.46557984, 189.38218125],
[114.46671081, 156.33419082, 139.91508329],
[ 60.50504249, 99.93731908, 100.45057895],
[152.42347775, 174.89832302, 153.38496535],
[ 24.28310161, 36.94729513, 65.34303884],
[223.21554659, 224.39480287, 222.18355735]
])

Como podemos ver, cada centroide es un array de punto flotante de 3 dimensiones. Por lo tanto, necesitamos convertir este array a un array de enteros.

k8_colors = np.uint8(k8_means.cluster_centers_)
k8_colors

array([
[ 42, 67, 84],
[ 87, 127, 119],
[191, 194, 189],
[114, 156, 139],
[ 60, 99, 100],
[152, 174, 153],
[ 24, 36, 65],
[223, 224, 222]
], dtype=uint8)

Grafiquemos los centros en el espacio de 3 dimensiones:

image.plot_rgb_vector(k8_colors, alpha=1, s=100, title="8-means centroids")

Para obtener la imagen cuantizada, simplemente reemplazamos el vector de features de cada píxel por el centroide del cluster al que pertenece:

k8_rgb_vector = k8_colors[k8_means.labels_]

Por último, podemos mostrar la imagen cuantizada:

k8_image = image.rgb_vector_to_image(k8_rgb_vector, img.shape)
k8_image.title = "8-mean Image"
image.display_images(k8_image)

Para acelerar el proceso creamos una función quantize_image() que, dada la imagen a cuantizar y el valor de \(k\), entrena el modelo \(k\)-means, muestra el estado del entrenamiento (tiempo e iteraciones ejecutadas) y despliega los centroides junto con la imagen cuantizada resultante. La función devuelve la imagen cuantizada.

def quantize_image(img: image.Image, k: int, random_state: int):
    kmeans = model.KMeans(n_clusters=k, random_state=random_state)

    # Training
    print("[Training]")
    %time kmeans.fit(img.rgb_vector)
    print(f"Iterations: {kmeans.iterations}")

    # Results
    print("\n[Results]")
    centroids = np.uint8(kmeans.cluster_centers_)
    print(f"{k}-means centroids:")
    image.plot_rgb_vector(centroids, alpha=1, s=100, title=f"{k}-means centroids")

    quantized_img = image.rgb_vector_to_image(centroids[kmeans.labels_], img.shape)
    quantized_img.title = f"{k}-mean image"

    print("Quantized image:")
    image.display_images(quantized_img)

    return quantized_img

\(16\)-means

Ahora podemos entrenar el algoritmo de \(16\)-means con la función que creamos:

k16_image = quantize_image(img, k=16, random_state=44)

[Training]
CPU times: total: 25.8 s
Wall time: 26.2 s
Iterations: 100

[Results]
16-means centroids:

Quantized image:

\(32\)-means

Finalmente, cuantizamos la imagen con \(k = 32\). Por lo que observamos antes, esperamos que la imagen cuantizada se parezca más a la original y que el tiempo de ejecución aumente.

k32_image = quantize_image(img, k=32, random_state=44)

[Training]
CPU times: total: 36.3 s
Wall time: 36.5 s
Iterations: 76

[Results]
32-means centroids:

Quantized image:

Comparación

Comparando los resultados obtenidos con \(k=8\), \(k=16\) y \(k=32\), podemos observar cómo la elección de la cantidad de centros influye en la apariencia visual de la imagen. Un valor de \(k\) más chico puede dar una reducción más marcada en la variedad de colores, mientras que un \(k\) más grande puede preservar detalles más finos en la distribución de color.

image.display_images(k8_image, k16_image, k32_image, figsize=(20, 15))

Vale la pena señalar que, además del impacto visual, la complejidad computacional del algoritmo aumenta con un valor de \(k\) más grande. Por lo tanto, la elección de \(k\) implica un compromiso entre eficiencia computacional y nivel de detalle preservado en la representación del color.

Medida de calidad y tiempo de ejecución

Podemos cuantificar la calidad de la imagen con una función de costo de los clusters, que mide la distancia euclídea al cuadrado entre cada vector de features \(x^{(i)}\) y su representante (es decir, el centroide del cluster) \(z_j\) del cluster \(C_j\).

\[\text{cost}(z_1, \dots, z_k) = \sum_{j=1}^k \sum_{i \in C_j} ||x^{(i)} - z_j||^2\]

def clusters_cost(X, centroids, labels):
    cost = 0
    for j in range(len(centroids)):
        cluster = X[np.where(labels == j)]
        distance = np.linalg.norm(cluster - centroids[j])**2
        cost += np.sum(distance)
    return cost

Si analizamos nuestro dataset \(X\), que representa los píxeles de la imagen, podemos contar la cantidad de colores distintos que tiene:

len(np.unique(img.rgb_vector, axis=0))

25786

Una cuantización con costo de clusters \(0\) requeriría un valor \(k=25786\), ya que necesitaríamos un centroide por cada color único de la imagen. Eso no es viable por la carga computacional que implica, y además la imagen cuantizada resultante no serviría de mucho. Por lo tanto, tenemos que limitar el rango de \(k\) para analizar cómo evolucionan el costo de los clusters y el tiempo de ejecución a medida que \(k\) aumenta.

La función model.plot_k_metrics() grafica cómo cambian el costo de los clusters y el tiempo de ejecución para distintos valores de \(k\). Para usarla, primero necesitamos entrenar el modelo para la lista de valores de \(k\) y registrar el costo de los clusters y el tiempo de ejecución en cada paso. Para eso, y para manejar el cómputo intensivo de forma más eficiente, creamos la función record_k_metrics(). Usamos la clase sklearn.cluster.KMeans, muy optimizada, para mejorar significativamente el tiempo de ejecución. Notemos que nuestra implementación, model.KMeans, tiene una interfaz idéntica a la de sklearn.cluster.KMeans, salvo por el parámetro n_init, que no está implementado en model.KMeans.

from sklearn.cluster import KMeans
from time import time

def record_k_metrics(
    img: image.Image, 
    k_values: list[int], 
    random_state: int,
    ):
    # Train the model for each k
    execution_time_list = []
    clusters_cost_list = []
    for k in k_values:
        kmeans = KMeans(n_clusters=k, random_state=random_state, n_init=5)
        # Execution time
        start_time = time()
        kmeans.fit(img.rgb_vector)
        end_time = time()
        execution_time_list.append(end_time - start_time)
        # Clusters cost
        centroids = np.uint8(kmeans.cluster_centers_)
        cost = clusters_cost(img.rgb_vector, centroids, kmeans.labels_)
        clusters_cost_list.append(cost)

    return clusters_cost_list, execution_time_list

Podemos visualizar el costo de los clusters y el tiempo de ejecución para \(k = 2, 4, 6, \dots, 64\):

k_values = list(range(2, 65, 2))
clusters_cost_list, execution_time_list = record_k_metrics(img, k_values, random_state=44)
model.plot_k_metrics(k_values, clusters_cost_list, execution_time_list)

Como \(k\)-means es un algoritmo estocástico, la relación entre \(k\) y las métricas introduce algo de ruido, lo que produce fluctuaciones ocasionales en la función. Aun así, a medida que \(k\) aumenta se hace evidente una tendencia decreciente en el costo de los clusters y una tendencia creciente en el tiempo de ejecución.

Conclusión

En esta exploración de la cuantización de imágenes con el algoritmo \(k\)-means nos metimos en los detalles de la reducción de color sobre una imagen de muestra de una flor. Variando la cantidad de clusters (\(k\)) y aplicando clustering \(k\)-means, observamos el impacto sobre la calidad de la imagen y la complejidad computacional.

Experimentar con \(k\)-means usando \(8\), \(16\) y \(32\) clusters nos dio una idea clara de los compromisos entre variedad de color y eficiencia computacional. Visualizamos las imágenes resultantes, analizamos los centroides de los clusters y comparamos la función de costo para distintos valores de k.

Además, exploramos la influencia de \(k\) sobre el tiempo de ejecución y el costo de los clusters. La función plot_k_metrics() nos dio una vista completa de estas métricas, iluminando la relación entre la cantidad de clusters y el rendimiento computacional.

Este proyecto, aunque es un comienzo humilde, marca un paso importante para entender los matices de la cuantización de imágenes. Tus comentarios y opiniones son muy valiosos mientras seguimos refinando y ampliando esta exploración de conceptos de machine learning.

¡Gracias por acompañarme en este recorrido! Espero con ganas tus ideas y sugerencias.