学習データとk近傍法を用いた文字認識

前提

kmnistという学習データとk近傍法を用いて
文字の認識を行う

実現したいこと

エラーを治し、適切にプログラムが作動する

発生している問題・エラーメッセージ

--------------------------------------------------------------------------- error Traceback (most recent call last) <ipython-input-45-0248c8fd3ca7> in <module> 1 knn = cv2.ml.KNearest_create() ----> 2 knn.train(train, cv2.ml.ROW_SAMPLE, train_labels) 3 ret, results, neighbours, dist = knn.findNearest(test, 5) error: OpenCV(4.6.0) /io/opencv/modules/ml/src/data.cpp:295: error: (-215:Assertion failed) (layout == ROW_SAMPLE && responses.rows == nsamples) || (layout == COL_SAMPLE && responses.cols == nsamples) in function 'setData'

該当のソースコード

python

1x_train = np.load('kmnist-train-imgs.npz')['arr_0']2x_test = np.load('kmnist-test-imgs.npz')['arr_0']3y_train = np.load('kmnist-train-labels.npz')['arr_0']4y_test = np.load('kmnist-test-labels.npz')['arr_0']5 6print(x_train.shape, x_test.shape, y_train.shape, y_test.shape)7 8fig,axes = plt.subplots(nrows=10, ncols=40, figsize=(12, 12))9for r in range(10):10 data = x_train[y_train==r]11 for c in range(40):12 axes[r, c].axis("off")13 axes[r, c].imshow(data[c], cmap='Greys_r')14plt.savefig('kmnist.png')15plt.show()16 17img = cv2.imread('kmnist.png')18 19x = np.array(data)20train = x[:,:50].reshape(-1,400).astype(np.float32) # Size = (2500,400)21test = x[:,50:100].reshape(-1,400).astype(np.float32) # Size = (2500,400)22 23k = np.arange(10)24train_labels = np.repeat(k,250)[:,np.newaxis]25test_labels = train_labels.copy()26 27knn = cv2.ml.KNearest_create()28knn.train(train, cv2.ml.ROW_SAMPLE, train_labels)29ret, results, neighbours, dist = knn.findNearest(test, 5)30

試したこと

以下のサイトを参考にし試行錯誤しましたがエラーは治らず非常に困っています。
https://stackoverflow.com/questions/49994760/opencv-assertion-failed-in-function-setdata