Lưu hình ảnh phân loại đúng với CNN

[Hoang_The_Nam (Hoàng Thế Nam)]

Mình đang tìm hiểu về nhận dạng chữ viết tay MNIST với CNN. có 1 yêu cầu là : trong tập ảnh test gồm 10000 ảnh , lưu lại 1000 ảnh được nhận dạng ( phân lớp ) đúng xuống thư mục , tức mỗi label 100 ảnh ( từ 0 => 9 ) . có ai làm được xin chỉ giáo với ạ ! tks
đây là code :

import keras
from keras.datasets import mnist 
from keras.layers import Dense, Activation, Flatten, Conv2D, MaxPooling2D
from keras.models import Sequential
from keras.utils import to_categorical
import numpy as np
import matplotlib.pyplot as plt

(train_X,train_Y), (test_X,test_Y) = mnist.load_data()

train_X = train_X.reshape(-1, 28,28, 1)
test_X = test_X.reshape(-1, 28,28, 1)

train_X = train_X.astype('float32')
test_X = test_X.astype('float32')
train_X = train_X / 255
test_X = test_X / 255

train_Y_one_hot = to_categorical(train_Y)
test_Y_one_hot = to_categorical(test_Y)

model = Sequential()

model.add(Conv2D(64, (3,3), input_shape=(28, 28, 1)))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(2,2)))

model.add(Conv2D(64, (3,3)))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(2,2)))

model.add(Flatten())
model.add(Dense(64))

model.add(Dense(10))
model.add(Activation('softmax'))

model.compile(loss=keras.losses.categorical_crossentropy, 
optimizer=keras.optimizers.Adam(),metrics=['accuracy'])

model.fit(train_X, train_Y_one_hot, batch_size=64, epochs=1)

test_loss, test_acc = model.evaluate(test_X, test_Y_one_hot)
print('Test loss', test_loss)
print('Test accuracy', test_acc)
model.save('123.model')
predictions = model.predict(test_X)
print(np.argmax(np.round(predictions[235])))

plt.imshow(test_X[235].reshape(28, 28), cmap = 'Greys_r')
plt.show()