import tensorflow
from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Flatten
from tensorflow.keras.layers import Conv2D, MaxPooling2D
import matplotlib.pyplot as plt

(trainX, trainY), (testX, testY) = mnist.load_data()

img_rows, img_cols = 28, 28

trainX = trainX.reshape(trainX.shape[0], img_rows, img_cols, 1)
testX = testX.reshape(testX.shape[0], img_rows, img_cols, 1)
input_shape = (img_rows, img_cols, 1) # 1 karena grayscale, 3 untuk berwarna

trainX = trainX.astype('float32')
testX = testX.astype('float32')

trainX /= 255
testX /= 255

trainY = tensorflow.keras.utils.to_categorical(trainY, 10)
testY = tensorflow.keras.utils.to_categorical(testY, 10)

model = Sequential()
model.add(Conv2D(32, kernel_size=(3,3), activation='relu', input_shape = input_shape))
model.add(Conv2D(64, (3,3), activation = 'relu'))
model.add(MaxPooling2D(pool_size=(2,2)))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dense(10, activation='softmax'))

model.compile(loss=tensorflow.keras.losses.categorical_crossentropy, optimizer="adam", metrics=['accuracy'])
model.fit(trainX, trainY, batch_size=128, epochs=12, verbose=2)
print (model.summary())
model.save("mnist.h5")