classify.py 文件源码

def test_models(models_dir, test_dir, out_dir):
    models_dir = utils.abs_path_dir(models_dir) + "/"
    test_dir = utils.abs_path_dir(test_dir) + "/"
    utils.create_dir(out_dir)
    test_files = os.listdir(test_dir)
    models = os.listdir(models_dir)
    for model in models:
        utils.print_success(model)
        pred_dir = out_dir + model + "/"
        utils.create_dir(pred_dir)
        clf = joblib.load(models_dir + model + "/" + model + ".pkl")
        for index, test_file in enumerate(test_files):
            print(str(index) + "\t" + test_file)
            sys.stdout.write("\033[F")
            sys.stdout.write("\033[K")
            test_features = read_test_file(test_dir + test_file)
            predictions = clf.predict_proba(test_features)
            with open(pred_dir + test_file, "w") as filep:
                for pred in predictions:
                    filep.write(str(pred[0]) + "\n")
        sys.stdout.write("\033[K")