knitting

knitting.py

"""
Author: Corey Lynch
Date: 10/4/12

example: python knitting.py 4 test.jpg 300 200
"""

import numpy as np
import pylab as pl
from sklearn.cluster import KMeans
from sklearn.metrics import euclidean_distances
from sklearn.datasets import load_sample_image
from sklearn.utils import shuffle
from time import time
from PIL import Image
import sys

# input
n_colors = int(sys.argv[1])
image_on_disk = sys.argv[2]


image_loaded = Image.open(image_on_disk)
width = int(sys.argv[3])
height = int(sys.argv[4])
image_downsample = image_loaded.resize((width, height), Image.ANTIALIAS)

image_to_knit = np.array(image_downsample, dtype=np.float64) / 255
w, h, d = original_shape = tuple(image_to_knit.shape)
assert d == 3
image_array = np.reshape(image_to_knit, (w * h, d))

print "Fitting estimator on a small sub-sample of the data"
t0 = time()
print image_array.shape
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)
rows = zip(*([iter(labels)]*599))

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()
dist = euclidean_distances(codebook_random, image_array, squared=True)
labels_random = dist.argmin(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

pl.figure(1)
pl.clf()
ax = pl.axes([0, 0, 1, 1])
pl.axis('off')
pl.title('Original image (96,615 colors)')
pl.imshow(image_to_knit)

pl.figure(2)
pl.clf()
ax = pl.axes([0, 0, 1, 1])
pl.axis('off')
pl.title('Quantized image (%s colors, K-Means)' % int(n_colors))
new_image = recreate_image(kmeans.cluster_centers_, labels, w, h)
pl.imshow(new_image)
pl.show()


# new_image == 400 * 599 numpy array with