from os import listdir
from os.path import basename, join
import numpy as np
from PIL import Image
from sklearn import svm
from sklearn.model_selection import learning_curve, ShuffleSplit
import matplotlib.pyplot as plt

# 图像尺寸
width, height = 30, 60
def loadDigits(dstDir='datasets'):
    # 获取所有图像文件名
    digitsFile = [join(dstDir,fn) for fn in listdir(dstDir)
                  if fn.endswith('.jpg')]
    # 按编号排序
    digitsFile.sort(key=lambda fn: int(basename(fn)[:-4]))
    # digitsData用于存放读取的图片中数字信息
    # 每个图片中所有像素值存放于digitsData中的一行数据
    digitsData = []
    for fn in digitsFile:
        with Image.open(fn) as im:
            # 原始图像尺寸不一样，先调整成统一尺寸，确保每个子列表的长度相同
            digitsData.append(np.array(im.resize((width,height)))
                              .mean(axis=2).flatten('F').tolist())
    # digitsLabel用于存放图片中数字的标准分类
    with open(fr'{dstDir}\digits.txt') as fp:
        digitsLabel = fp.read().splitlines(keepends=False)
    return (digitsData, digitsLabel)

# 加载数据
X, y = loadDigits()
print('数据加载完成。')

svcClassifier = svm.SVC(kernel='linear', C=1000, gamma=0.001)
# 使用20%的样本进行测试，80%的样本进行训练，重复10次，计算每次的得分
cv = ShuffleSplit(test_size=0.2, train_size=0.8, n_splits=10)
# 使用交叉验证计算不同数量样本作为训练集时的模型平均得分
train_sizes, train_scores, test_scores = learning_curve(svcClassifier, X, y, cv=cv,
                                                        train_sizes=np.linspace(0.1,1.0,10))
# 对于不同大小的训练集10次训练与测试结果得分求均值
train_scores = np.mean(train_scores, axis=1)
test_scores = np.mean(test_scores, axis=1)

plt.plot(train_sizes, train_scores, 'r-o', lw=2, label='training score')
plt.plot(train_sizes, test_scores, 'g--v', lw=2, label='cross validation score')
plt.legend()

plt.show()