try:
  # Colab only
  %tensorflow_version 2.x
except Exception:
  pass

import tensorflow as tf
from tensorflow import keras

import numpy as np
import matplotlib.pyplot as plt
import collections as col

mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

print(train_images.shape)
print(train_labels.shape)
print(test_images.shape)
print(test_labels.shape)

col.Counter(train_labels)
col.Counter(test_labels)
plt.hist(train_labels, color = 'green', alpha = .5)
plt.hist(test_labels, color = 'green', alpha = .5)

plt.figure(figsize = (10, 10))        # 전체 피겨 사이즈를 10*10으로
for i in range(25) :
  plt.subplot(5, 5, i + 1)
  plt.xticks([])                      # x 눈금 제거
  plt.yticks([])                      # y 눈금 제거
  plt.xlabel(train_labels[i])         # x 축 라벨에 정답
  plt.imshow(train_images[i])

col.Counter(train_images[0].reshape(784))    # 0 번째 이미지(24*24 2차원 배열)를 784사이즈의 1차원 배열로 reshape하여 요소(색을 결정하는 수치)를 확인해 보았다.

train_images = train_images / 255.0
test_images = test_images / 255.0

col.Counter(train_images[0].reshape(784))

plt.figure(figsize = (10, 10))
plt.imshow(train_images[0])
plt.colorbar()
plt.show()

model = keras.Sequential()
model.add(keras.layers.Flatten(input_shape = (28, 28)))
model.add(keras.layers.Dense(128, activation = 'relu'))
model.add(keras.layers.Dense(10, activation = 'softmax'))

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

model.fit(train_images, train_labels, epochs=10)

loss, accuracy = model.evaluate(test_images,  test_labels, verbose = 1)
print('\ntest loss\t: ', loss)
print('test accuracy\t: ', accuracy)

predictions = model.predict(test_images)
plt.figure(figsize=(20,20))
for i in range(25) :
  plt.subplot(5, 5, i+1)
  plt.imshow(test_images[i])
  plt.xticks([])
  plt.yticks([])
  description = str(np.argmax(predictions[i])) + " / " + str(test_labels[i])
  plt.xlabel(description)
plt.show()