Commit 6ff538c4f163d34a30a3bf542ad0f9cca2549cbe
1 parent
6c40a57b2c
Exists in
master
Allow you to test if kmeans work well.
Showing 1 changed file with 34 additions and 0 deletions Side-by-side Diff
bin/test-kmeans.py
| 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) |