From 993ea26cfef22ea3fc8abf3c6973de1d0fc297c0 Mon Sep 17 00:00:00 2001 From: Mathias Quillot Date: Mon, 15 Jul 2019 09:35:14 +0200 Subject: [PATCH] take into account list option --- bin/extract_kmeans.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/bin/extract_kmeans.py b/bin/extract_kmeans.py index 2a95713..a216527 100644 --- a/bin/extract_kmeans.py +++ b/bin/extract_kmeans.py @@ -13,7 +13,7 @@ import sys parser = argparse.ArgumentParser(description="extract clusters") parser.add_argument("model", type=str, help="k-means model pickle") parser.add_argument("features", type=str, help="features") -parser.add_argument("list", type=str, help="list file") +parser.add_argument("--list", type=str, default=None, help="list file") parser.add_argument("--outfile", type=str, default=None, help="output file std") args = vars(parser.parse_args()) @@ -31,13 +31,16 @@ else: features = read_file(FEATURES) feat_ind = index_by_id(features) -lst = read_file(LST) +if LST is not None: + lst = read_file(LST) +else: + lst = features kmeans = pickle.load(open(MODEL, "rb")) - # -- CONVERT TO NUMPY -- X = np.asarray([feat_ind[x[0][0]][x[0][3]][1] for x in lst]) + predictions = kmeans.predict(X) for i, line in enumerate(lst): -- 1.8.2.3