Commit 6ff538c4f163d34a30a3bf542ad0f9cca2549cbe

Authored by Mathias Quillot
1 parent 6c40a57b2c
Exists in master

Allow you to test if kmeans work well.

Showing 1 changed file with 34 additions and 0 deletions Inline Diff

File was created 1 '''
2 '''
3
4
5 import argparse
6 import numpy as np
7 from sklearn.cluster import KMeans
8 from data import read_file, index_by_id
9 import pickle
10
11 parser = argparse.ArgumentParser(description="...")
12 parser.add_argument("kmeans", type=str, help="kmean saved file")
13 parser.add_argument("features", type=str, help="features file")
14 parser.add_argument("lst", type=str, help="lst file")
15
16 args = parser.parse_args()
17 KMEAN_FILE = args.kmeans
18 FEATURES_FILE = args.features
19 LST_FILE = args.lst
20
21 # Load features and lst
22 features = read_file(FEATURES_FILE)
23 features_ind = index_by_id(features)
24
25 lst = read_file(LST_FILE)
26
27 # Load Kmeans
28 kmeans = pickle.load(open(KMEAN_FILE, "rb"))
29
30 # Get all x
31 X = np.asarray([features_ind[x[0][0]][x[0][3]][1] for x in lst])
32 predicts = kmeans.predict(X)
33 print(np.unique(predicts).shape)
34
35