Clustering example

example/test_clustering.py

@@ -0,0 +1,50 @@
+import sys
+import os
+sys.path.append("..")  # Adds higher directory to python modules path.
+from img_to_vec import Img2Vec
+from PIL import Image
+import numpy as np
+from sklearn.cluster import KMeans
+from sklearn.decomposition import PCA
+
+input_path = './test_images'
+files = os.listdir(input_path)
+
+img2vec = Img2Vec()
+vec_length = 512  # Using resnet-18 as default
+
+samples = 5  # Amount of samples to take from input path
+k_value = 2  # How many clusters
+
+vec_mat = np.zeros((samples, vec_length))
+sample_indices = np.random.choice(range(0, len(files)), size=samples, replace=False)
+
+print('Reading images...')
+for index, i in enumerate(sample_indices):
+    file = files[i]
+    filename = os.fsdecode(file)
+    img = Image.open(os.path.join(input_path, filename))
+    vec = img2vec.get_vec(img)
+    vec_mat[index, :] = vec
+
+print('Applying PCA...')
+reduced_data = PCA(n_components=2).fit_transform(vec_mat)
+kmeans = KMeans(init='k-means++', n_clusters=k_value, n_init=10)
+kmeans.fit(reduced_data)
+
+for i in set(kmeans.labels_):
+    try:
+        os.mkdir('./' + str(i))
+    except FileExistsError:
+        continue
+
+print('Predicting...')
+preds = kmeans.predict(reduced_data)
+
+print('Moving images...')
+for index, i in enumerate(sample_indices):
+    file = files[i]
+    filename = os.fsdecode(file)
+    os.rename(input_path + '/' + filename, './' + str(preds[index]) + '/' + filename)
+
+print('Done!')

example/test_visualize.py

@@ -0,0 +1,50 @@
+import sys
+import os
+sys.path.append("..")  # Adds higher directory to python modules path.
+from img_to_vec import Img2Vec
+from PIL import Image
+from sklearn.manifold import TSNE
+import numpy as np
+from ggplot import *
+import pandas as pd
+
+input_path = './test_images'
+files = os.listdir(input_path)
+
+img2vec = Img2Vec()
+vec_length = 512  # Using resnet-18 as default
+
+samples = 50  # Amount of samples to take from input path
+
+vec_mat = np.zeros((samples, vec_length))
+sample_indices = np.random.choice(range(0, len(files)), size=samples, replace=False)
+
+# For each test sample, we store the label and append the img vector to our matrix
+labels = []
+for index, i in enumerate(sample_indices):
+    file = files[i]
+    filename = os.fsdecode(file)
+    img = Image.open(os.path.join(input_path, filename))
+    vec = img2vec.get_vec(img)
+    vec_mat[index, :] = vec
+    if 'dog' in filename:
+        labels.append('dog')
+    elif 'cat' in filename:
+        labels.append('cat')
+    elif 'face' in filename:
+        labels.append('face')
+
+t_sne = TSNE()
+
+tsne_output = t_sne.fit_transform(vec_mat)
+
+df_tsne = pd.DataFrame(columns=['x-tsne', 'y-tsne', 'label'])
+df_tsne['x-tsne'] = tsne_output[:, 0]
+df_tsne['y-tsne'] = tsne_output[:, 1]
+df_tsne['label'] = labels
+
+chart = ggplot(df_tsne, aes(x='x-tsne', y='y-tsne', color='label')) \
+        + geom_point(size=90, alpha=0.8) \
+        + ggtitle("tSNE dimensions colored by image label")
+
+print(chart)

img_to_vec.py

@@ -8,9 +8,11 @@
 class Img2Vec():
 
     def __init__(self, cuda=False, model='resnet-18', layer='default', layer_output_size=512):
-        """
-            Parameters:
-                see docs: https://github.com/christiansafka/img2vec
+        """ Img2Vec
+        :param cuda: If set to True, will run forward pass on GPU
+        :param model: String name of requested model
+        :param layer: String or Int depending on model.  See more docs: https://github.com/christiansafka/img2vec.git
+        :param layer_output_size: Int depicting the output size of the requested layer
         """
         self.cuda = cuda
         self.layer_output_size = layer_output_size
@@ -26,12 +28,11 @@ def __init__(self, cuda=False, model='resnet-18', layer='default', layer_output_
                                               std=[0.229, 0.224, 0.225])
         self.to_tensor = transforms.ToTensor()
 
-    def get_vec(self, img):
-        """
-            Parameters:
-                img: PIL image
-            Returns:
-                Numpy vector representing input image
+    def get_vec(self, img, tensor=False):
+        """ Get vector embedding from PIL image
+        :param img: PIL Image
+        :param tensor: If True, get_vec will return a FloatTensor instead of Numpy array
+        :returns: Numpy ndarray
         """
         if self.cuda:
             image = Variable(self.normalize(self.to_tensor(self.scaler(img))).unsqueeze(0)).cuda()
@@ -47,9 +48,17 @@ def copy_data(m, i, o):
         h_x = self.model(image)
         h.remove()
 
-        return my_embedding.numpy()
+        if tensor:
+            return my_embedding
+        else:
+            return my_embedding.numpy()
 
     def _get_model_and_layer(self, model_name, layer):
+        """ Internal method for getting layer from model
+        :param model_name: model name such as 'resnet-18'
+        :param layer: layer as a string for resnet-18 or int for alexnet
+        :returns: pytorch model, selected layer
+        """
         if model_name == 'resnet-18':
             model = models.resnet18(pretrained=True)
             if layer == 'default':