fitting_kmeans_image.py

from dask_ml.cluster import KMeans
from PIL import Image
import numpy as np

#given the fitted k-means model and a vector of pixels, returns the new value for each pixel.
def clustered_pixels(x_fit, pixels):
    labels = x_fit.predict(pixels).compute()
    res= x_fit.cluster_centers_[labels] 
    return res

#defined in the previous gist.
np_im = pixels_from_path(file_path)
pixels = vector_of_pixels(np_im)

km = KMeans(n_clusters=k) # k = how many colors we want the new image to have
x_fit = km.fit(pixels) #we fit the centroids and labels.
clust_pixels = clustered_pixels(x_fit, pixels)

pic = reshape_pixels(clust_pixels, width, height) #turn vector to matrix again.