Blame view

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