KNN 图像分类入门

"""KNN 图像分类入门 —— 站内原创示例代码

用途：用 scikit-learn 自带的 8x8 手写数字图像数据集，训练 K 近邻分类器并评估准确率。
      这是理解「图像分类」最基本流程的最小示例：数据 -> 特征 -> 训练 -> 评估。
依赖：scikit-learn（pip install scikit-learn）
运行：python knn-classify.py
许可：MIT（开源拾遗原创，可自由使用与修改）
"""

from sklearn.datasets import load_digits
from sklearn.metrics import accuracy_score, confusion_matrix
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

K = 5


def main():
    # 1. 载入数据：1797 张 8x8 手写数字图片，每张展平成 64 维特征
    digits = load_digits()
    images = digits.images
    features = digits.data
    labels = digits.target

    print("样本数:", features.shape[0], " 特征维度:", features.shape[1])

    # 2. 划分训练集与测试集
    x_train, x_test, y_train, y_test = train_test_split(
        features, labels, test_size=0.2, random_state=42, stratify=labels
    )

    # 3. 训练 KNN（直接比较像素距离，是最朴素的图像分类基线）
    model = KNeighborsClassifier(n_neighbors=K)
    model.fit(x_train, y_train)

    # 4. 评估
    predictions = model.predict(x_test)
    accuracy = accuracy_score(y_test, predictions)
    print("K =", K)
    print("测试集准确率: %.4f" % accuracy)

    # 5. 看一眼混淆情况
    matrix = confusion_matrix(y_test, predictions)
    print("混淆矩阵对角线（每类正确数）:", [int(matrix[i][i]) for i in range(10)])

    # 自检：KNN 在该数据集上准确率应明显高于 0.9
    assert accuracy > 0.9, accuracy

    # 展示第一张测试图的预测结果
    print("第一张测试图真实值:", int(y_test[0]), " 预测值:", int(predictions[0]))
    print("图片像素（8x8）:")
    for row in x_test[0].reshape(8, 8):
        print("".join("#" if value > 8 else "." for value in row))
    print("断言通过")


if __name__ == "__main__":
    main()