# 构建手写数字OCR的应用程序
#
# 需要一些 train_data 和 test_data。 OpenCV 附带一个图像digits.png，其中有 5000 个手写数字（每个数字 500 个）。每个数字都是一个 20x20 的图像。、
# 第一步是将这个图像分成 5000 个不同的数字。对于每个数字，我们将其展平为 400 像素的一行。那就是我们的特征集，即所有像素的强度值。这是我们可以创建的最简单的功能集。
# 第二部使用每个数字的前 250 个样本作为 train_data，剩下的250 个样本作为 test_data
# 然后训练数据；

import numpy as np
import cv2

img = cv2.imread('digits.png')  # 原始图像包括5000个不同的数字
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# 分割图像为5000个单元，每个20x20像素
cells = [np.hsplit(row, 100) for row in np.vsplit(gray, 50)]

# 转为Numpy数组，shape将为 (50,100,20,20)
x = np.array(cells)

# 准备训练集和测试集
train = x[:, :50].reshape(-1, 400).astype(np.float32)  # Size = (2500,400)
test = x[:, 50:100].reshape(-1, 400).astype(np.float32)  # Size = (2500,400)

# 为训练数据和测试数据创建标签
k = np.arange(10)
train_labels = np.repeat(k, 250)[:, np.newaxis]
test_labels = train_labels.copy()

# 初始化KNN，训练模型，然后对测试集使用k=1进行测试
knn = cv2.ml.KNearest_create()
#print('knn: ', knn)
knn.train(train, cv2.ml.ROW_SAMPLE, train_labels)
ret, result, neighbours, dist = knn.findNearest(test, k=5)

# 检查分类的准确性
# 将预测分类的结果与其所属的测试标签进行比较，判断成功还是失败
matches = result == test_labels
correct = np.count_nonzero(matches)
accuracy = correct * 100.0 / result.size

# 识别手写数字得到了91.76% 的准确率。提高准确性的一种选择是添加更多用于训练的数据，尤其是错误的数据。
# 可以把训练的数据/模型保存下来，这样下次就能直接从文件中加载数据/模型，然后进行分类，可借助Numpy函数（如 np.savetxt、np.savez、np.load 等）来保存模型
print("OCR digits accuracy: ", accuracy)


image = cv2.imread('9/9.jpg',0)
#img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY);
image.resize(20,20)
ret, image = cv2.threshold(image, 0, 255, cv2.THRESH_OTSU | cv2.THRESH_BINARY_INV)
testImage = image.reshape(-1,400).astype(np.float32)

# 用所有数据作为训练数据

ret,result,neighbours,dist = knn.findNearest(testImage, 1)
print(result)