sklearn的简单使用

sklearn的简单使用

import numpy as np
from sklearn import neighbors
from sklearn.model_selection import train_test_split
import joblib

import os


# 将图片数据转化
def image2vector(filename):
    returnVect = np.zeros(1024)
    fr = open(filename)
    for i in range(32):
        lineStr = fr.readline()
        for j in range(32):
            returnVect[i * 32 + j] = int(lineStr[j])
    return returnVect


# 获取本地数据
def getSourceDatas(filepath):
    datas = []
    labels = []
    for files in os.listdir(filepath):
        datas.append(image2vector(filepath + "/" + files))
        labels.append(files.strip().split("_")[0])
    return datas, labels


if __name__ == "__main__":
    # 获取数据
    dir_path = os.getcwd()
    train_path = dir_path + "/trainingDigits"
    test_path = dir_path + "/testDigits"
    train_datas, train_labels = getSourceDatas(train_path)
    test_datas, test_labels = getSourceDatas(test_path)
    x_train, x_test, y_train, y_test = train_test_split(train_datas, train_labels, test_size=0.3)
    # 训练
    clf = neighbors.KNeighborsClassifier()
    clf.fit(train_datas, train_labels)
    # 获取测试数据正确率
    score = clf.score(test_datas, test_labels)
    print("正确率：" + str(score))
    # 保存模型
    joblib.dump(clf, 'clf.pkl')
    # 加载模型
    clf3 = joblib.load('clf.pkl')
    print("预测结果：")
    print(clf3.predict(test_datas[0:1]))
    print("test结果：")
    print(test_labels[0:1])

测试和训练数据以及源码