import numpy as np
import cv2 as cv


if __name__ == '__main__':
    
    img = cv.imread('/root/gitdata/zhou_test/KNN_test/digits.png')

    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    # Now we split the image to 5000 cells, each 20x20 size
    cells = [np.hsplit(row, 100) for row in np.vsplit(gray, 50)]
    # Make it into a Numpy array. It size will be (50, 100, 20, 20)
    x = np.array(cells)
    # Now we prepare train_data and test_data.
    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)
    # Create labels for train and test data
    k = np.arange(10)
    train_labels = np.repeat(k, 250)[:, np.newaxis]    # train_labels.shape=(2500, 1)
    test_labels = train_labels.copy()
    ml_traindata = cv.ml.TrainData_create(train, cv.ml.ROW_SAMPLE, train_labels)      # <class 'cv2.ml_TrainData'>
    # Initiate kNN, train the data, then test it with test data for k=1
    knn = cv.ml.KNearest_create()       # <class 'cv2.ml_KNearest'>
    # samples	training samples
    # layout	See ml::SampleTypes.
    # responses	vector of responses associated with the training samples.
    # retval = cv.ml_StatModel.train(samples, layout, responses)
    train_ret = knn.train(ml_traindata)
    # train_ret = knn.train(train, cv.ml.ROW_SAMPLE, train_labels)    # train_ret = True
    ret, result, neighbours, dist = knn.findNearest(test, k=5)
    tt2 = ml_traindata.getTestResponses()
    # result.shape = (2500, 1),  neighbours.shape=(2500, 5),    dist.shape=(2500, 5)
    # Now we check the accuracy of classification
    # For that, compare the result with test_labels and check which are wrong
    matches = result == test_labels
    correct = np.count_nonzero(matches)
    accuracy = correct*100.0/result.size
    print(accuracy)

import numpy as np
import cv2 as cv


if __name__ == '__main__':
    
    img = cv.imread('/root/gitdata/zhou_test/KNN_test/digits.png')

    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    # Now we split the image to 5000 cells, each 20x20 size
    cells = [np.hsplit(row, 100) for row in np.vsplit(gray, 50)]
    # Make it into a Numpy array. It size will be (50, 100, 20, 20)
    x = np.array(cells)
    # Now we prepare train_data and test_data.
    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)
    # Create labels for train and test data
    k = np.arange(10)
    train_labels = np.repeat(k, 250)[:, np.newaxis]    # train_labels.shape=(2500, 1)
    test_labels = train_labels.copy()
    ml_traindata = cv.ml.TrainData_create(train, cv.ml.ROW_SAMPLE, train_labels)      # <class 'cv2.ml_TrainData'>
    # Initiate kNN, train the data, then test it with test data for k=1
    knn = cv.ml.KNearest_create()       # <class 'cv2.ml_KNearest'>
    # samples	training samples
    # layout	See ml::SampleTypes.
    # responses	vector of responses associated with the training samples.
    # retval = cv.ml_StatModel.train(samples, layout, responses)
    train_ret = knn.train(ml_traindata)
    # train_ret = knn.train(train, cv.ml.ROW_SAMPLE, train_labels)    # train_ret = True
    ret, result, neighbours, dist = knn.findNearest(test, k=5)
    tt2 = ml_traindata.getTestResponses()
    # result.shape = (2500, 1),  neighbours.shape=(2500, 5),    dist.shape=(2500, 5)
    # Now we check the accuracy of classification
    # For that, compare the result with test_labels and check which are wrong
    matches = result == test_labels
    correct = np.count_nonzero(matches)
    accuracy = correct*100.0/result.size
    print(accuracy)

# Now we split the image to 5000 cells, each 20x20 size # Make it into a Numpy array. It size will be (50, 100, 20, 20) # Now we prepare train_data and test_data. # Size = (2500, 400) # Size = (2500, 400) # Create labels for train and test data # train_labels.shape=(2500, 1) # <class 'cv2.ml_TrainData'> # Initiate kNN, train the data, then test it with test data for k=1 # <class 'cv2.ml_KNearest'> # samples	training samples # layout	See ml::SampleTypes. # responses	vector of responses associated with the training samples. # retval = cv.ml_StatModel.train(samples, layout, responses) # train_ret = knn.train(train, cv.ml.ROW_SAMPLE, train_labels)    # train_ret = True # result.shape = (2500, 1),  neighbours.shape=(2500, 5),    dist.shape=(2500, 5) # Now we check the accuracy of classification # For that, compare the result with test_labels and check which are wrong