Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- # Authors: Robert Layton <robertlayton@gmail.com>
- # Olivier Grisel <olivier.grisel@ensta.org>
- # Mathieu Blondel <mathieu@mblondel.org>
- # http://scikit-learn.org/stable/auto_examples/cluster/plot_color_quantization.html
- #
- # License: BSD 3 clause
- from PIL import Image
- print(__doc__)
- import numpy as np
- import matplotlib.pyplot as plt
- from sklearn.cluster import KMeans
- from sklearn.metrics import pairwise_distances_argmin
- from sklearn.datasets import load_sample_image
- from sklearn.utils import shuffle
- from time import time
- n_colors = 4
- # Load the Summer Palace photo
- # image = load_sample_image("image.jpg")
- image = Image.open("lena.jpg")
- # Convert to floats instead of the default 8 bits integer coding. Dividing by
- # 255 is important so that plt.imshow behaves works well on float data (need to
- # be in the range [0-1])
- image = np.array(image, dtype=np.float64) / 255
- # Load Image and transform to a 2D numpy array.
- w, h, d = original_shape = tuple(image.shape)
- assert d == 3
- image_array = np.reshape(image, (w * h, d))
- print("Fitting model on a small sub-sample of the data")
- t0 = time()
- image_array_sample = shuffle(image_array, random_state=0)[:1000]
- kmeans = KMeans(n_clusters=n_colors, random_state=0).fit(image_array_sample)
- print("done in %0.3fs." % (time() - t0))
- # Get labels for all points
- print("Predicting color indices on the full image (k-means)")
- t0 = time()
- labels = kmeans.predict(image_array)
- print("done in %0.3fs." % (time() - t0))
- codebook_random = shuffle(image_array, random_state=0)[:n_colors + 1]
- print("Predicting color indices on the full image (random)")
- t0 = time()
- labels_random = pairwise_distances_argmin(codebook_random,
- image_array,
- axis=0)
- print("done in %0.3fs." % (time() - t0))
- def recreate_image(codebook, labels, w, h):
- """Recreate the (compressed) image from the code book & labels"""
- d = codebook.shape[1]
- image = np.zeros((w, h, d))
- label_idx = 0
- for i in range(w):
- for j in range(h):
- image[i][j] = codebook[labels[label_idx]]
- label_idx += 1
- return image
- # Display all results, alongside original image
- plt.figure(1)
- plt.clf()
- ax = plt.axes([0, 0, 1, 1])
- plt.axis('off')
- plt.title('Original image (96,615 colors)')
- plt.imshow(image)
- plt.figure(2)
- plt.clf()
- ax = plt.axes([0, 0, 1, 1])
- plt.axis('off')
- plt.title('Quantized image (64 colors, K-Means)')
- plt.imshow(recreate_image(kmeans.cluster_centers_, labels, w, h))
- plt.figure(3)
- plt.clf()
- ax = plt.axes([0, 0, 1, 1])
- plt.axis('off')
- plt.title('Quantized image (64 colors, Random)')
- plt.imshow(recreate_image(codebook_random, labels_random, w, h))
- plt.show()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement