import pickle
from os.path import join
from data import load_emb_data
from identification import IdentifyModel
# def train(emb_data):
# ide = IdentifyModel()
# ide.fit(emb_data)
# return ide
if __name__=='__main__':
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("--data-dir", help="data-dir")
ap.add_argument("--model-path", help="model-path")
ap.add_argument("--idx2path", help="idx2path")
ap.add_argument("--vector-dir", help="vector_dir")
args= vars(ap.parse_args())
name2file, name2emb = load_emb_data(args['data_dir'], vector_dir=args['vector_dir'], force=True)
idx2path = {name: [join(args['data_dir'], name, f) for f in files] for name, files in name2file.items()}
emb_data = name2emb
with open(args['idx2path'], 'wb') as f:
pickle.dump(idx2path, f)
ide = IdentifyModel()
ide.fit(emb_data)
ide.dump_classify_model(args['model_path'])
ap.add_argument("--data-dir", help="data-dir")
ap.add_argument("--known-vector-dir", help="known-vector-dir")
ap.add_argument("--ver-vector-dir", help="ver-vector-dir")
ap.add_argument("--model-path", help="model-path")
ap.add_argument("--idx2path", help="idx2path")
ap.add_argument("--threshold", type=float,help="threshold")
ap.add_argument("--k", type=int, help="n_top_candidate")
ap.add_argument("--output", help="output")
ap.add_argument("--batch-size", type=int, help="batch-size")
ap.add_argument("--tree-path", help="tree-path")
args= vars(ap.parse_args())
tree = Tree()
# if not exists(args['tree_path']):
start = time()
name2file, emb_data = load_emb_data(args['data_dir'], vector_dir=args['ver_vector_dir'])
data = {name: [join(args['data_dir'], name, f) for f in files] for name, files in name2file.items()}
print('loaded data: ', time() - start)
for name, paths in data.items():
_emb = emb_data[name][0]
test_img = TestImage(paths[0], _emb)
tree.append(test_img)
# pdb.set_trace()
ide_model = IdentifyModel()
ide_model.load_classify_model(args['model_path'])
# ide_model.known_vector_dir(args['known_vector_dir'])
# ide_model.set_threshold(args['threshold'])
ide_model.set_n_top_candidate(args['k'])
ide_model.load_idx2path(args['idx2path'])
print('ide_model.idx2path: ', ide_model.idx2path)