DigitRecognizer_GUI/model2.py at master · Shane120283483/DigitRecognizer_GUI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
#!/usr/bin/python
# -*- coding:utf-8 -*-
import numpy as np
import os
from PIL import Image
from sklearn.neighbors import KNeighborsClassifier as knn
from sklearn.externals import joblib


'''将图片转换为向量'''


def image_vector(fname):
    im = Image.open("mnist_train/"+fname).convert('L')
    im = im.resize((28, 28))
    tmp = np.array(im)
    vector = tmp.ravel()           # 转换成1*784的向量
    return vector


def image_vector_2(fname):
    im = Image.open("mnist_test/"+fname).convert('L')
    im = im.resize((28, 28))
    tmp = np.array(im)
    vector = tmp.ravel()           # 转换成1*784的向量
    return vector


'''训练60000张图片'''


def split_data(paths):
    fn_list = os.listdir(paths)
    X = []
    y = []
    d0 = fn_list
    for i, name in enumerate(d0):
        y.append(name[0])  # 获取文件名的第一个字符，例如0_train_1.bmp,则得到数字标签0
        X.append(image_vector(name))  # 获取图片并利用函数转换为1*784向量
    return X, y


'''构建分类器'''


def knn_clf(X_train, y_train_label):
    classifier = knn()
    classifier.fit(X_train, y_train_label)
    return classifier


'''保存模型'''


def save_model(model, output_name):
    joblib.dump(model, output_name)


'''训练模型'''


X_train_1, y_train_label_1 = split_data("mnist_train")
clf = knn_clf(X_train_1, y_train_label_1)
save_model(clf, 'mnist_knn60000.m')


'''加载模型'''


def load_model(model):
    joblib.load(model)


'''获取测试集'''


def testdata(filename):
    X = []
    X.append(image_vector_2(filename))
    return X


'''测试模型'''


def test(test_sample):
    load_model('mnist_knn60000.m')
    test_result = clf.predict(testdata(test_sample))
    # print(test_result[0])
    # if test_sample[0] == test_result:
       # print("Bingo")
    # else:
       # print("Wrong")
    return test_result[0]


# test("8_test_1170.bmp")