python 机器学习（手写数字识别）

# 导包
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
import joblib
from collections import Counter

# 1. 定义函数 show_digit(idx), 用于查看: 数字图片.
def show_digit(idx):
    # idx: 行索引, 即: 要哪行的数据.
    # 1. 读取数据, 获取df对象.
    data = pd.read_csv('data/手写数字识别.csv')
    # 细节: 非法值校验.
    if idx < 0 or idx > len(data) - 1 :
        return
    # 2. 获取数据, 即: 特征 + 标签
    x = data.iloc[:, 1:]
    y = data.iloc[:, 0]
    # 3. 查看下数据集.
    print(f'x的维度: {x.shape}')           # (42000, 784)
    print(f'y的各分类数量: {Counter(y)}')   # Counter({1: 4684, 7: 4401, 3: 4351, 9: 4188, 2: 4177, 6: 4137, 0: 4132, 4: 4072, 8: 4063, 5: 3795})

    # 4. 获取具体的 某张图片, 即: 某行数据 => 样本数据.
    # step1: 把图片转成 28 * 28像素(二维数组)
    digit = x.iloc[idx].values.reshape(28, 28)
    # step2: 显示图片.
    plt.imshow(digit, cmap='gray')      # 灰色显示 => 灰度图
    # step3: 取消坐标显示.
    plt.axis('off')
    # step4: 显示图片
    plt.show()

# 2. 定义函数 train_model(), 用于训练模型.
def train_model():
    # 1. 读取数据, 获取df对象.
    data = pd.read_csv('data/手写数字识别.csv')
    # 2. 获取数据, 即: 特征 + 标签
    x = data.iloc[:, 1:]
    y = data.iloc[:, 0]
    # 3. 数据的预处理.
    # step1: x轴(像素点)的 归一化处理.
    x = x / 255
    # step2: 区分训练集和测试集.       stratify: 按照y的类别比例进行分割
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, stratify=y, random_state=21)
    # 4. 训练模型.
    estimator = KNeighborsClassifier(n_neighbors=3)
    estimator.fit(x_train, y_train)
    # 5. 模型预估, 评测正确率.
    my_score = estimator.score(x_test, y_test)
    print(f'模型预估正确率: {my_score}')       # 0.9657142857142857
    # 6. 保存模型.
    joblib.dump(estimator, 'model/knn.pth')

# 3. 定义use_model()函数, 用于: 测试模型.
def use_model():       # pytest
    # 1. 加载图片.
    img = plt.imread('data/demo.png')       # 28 * 28像素
    plt.imshow(img, cmap='gray')    # 灰度图
    plt.show()

    # 2. 加载模型.
    estimator = joblib.load('model/knn.pth')

    # 3. 预测图片.
    img = img.reshape(1, -1)    # 效果等价于: reshape(1, 784)
    y_test = estimator.predict(img)
    print(f'预测的数字是: {y_test}')

# 4. 在main函数中测试.
if __name__ == '__main__':
    # 显示数字图片.
    # show_digit(10)
    # show_digit(21)

    # 训练模型
    # train_model()

    # 测试模型
    use_model()