[불친절한 튜토리얼] Grad-CAM, TensorFlow, Keras, MNIST

input_shape = (28, 28, 1)  # 입력 데이터 하나의 크기 (높이, 너비, 채널)
num_classes = 10  # 분류 클래스 개수

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

import tensorflow as tf
from tensorflow import keras

(X_train, y_train), (X_test, y_test) = keras.datasets.mnist.load_data()

print(X_train.shape)
print(y_train.shape)
print(X_test.shape)
print(y_test.shape)

X_train = X_train[..., np.newaxis]
X_train = X_train.astype(np.float32)
X_train = X_train / 255.0

X_test = X_test[..., np.newaxis]
X_test = X_test.astype(np.float32)
X_test = X_test / 255.0

y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

def get_basic_cnn(input_shape=(28, 28, 1), num_classes=10):
    cnn = keras.models.Sequential([
        keras.layers.Conv2D(32, 3, padding='same', activation='relu', input_shape=input_shape),
        keras.layers.MaxPool2D(),
        keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
        keras.layers.MaxPool2D(),
        keras.layers.Flatten(),
        keras.layers.Dense(num_classes),
        keras.layers.Softmax(),
    ])
    
    return cnn

model = get_basic_cnn()
model.summary()

model.compile(loss='categorical_crossentropy', optimizer='adam', metrics='acc')

history = model.fit(X_train, y_train, validation_split=0.2, epochs=10, batch_size=64)

def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
    # 마지막 feature maps에 대한 모델의 미분을 계산하기 위해서
    # 마지막 feature maps과 모델을 최종 예측값을 동시에 출력하는 모델을 새롭게 정의한다.
    # 이것은 케라스 문법이기 때문에 익숙치 않을 수 있다. 
    #    - Model 함수의 첫 번째 인자는 입력 객체, 두 번째 인자는 출력 객체이다.
    #    - 마지막 feature maps과 최종 예측값 두 개를 출력한다.
    grad_model = tf.keras.models.Model(
        [model.inputs], [model.get_layer(last_conv_layer_name).output, model.output]
    )

    # 이건 텐서플로우 문법이다.
    # 모델에 이미지가 입력되고, 마지막 feature maps과 최종 예측값을 구하는 일련의 연산 과정을
    # Computational graph라는 형태로 저장해놓는다. 이렇게 해야 미분값을 계산할 수 있다.
    with tf.GradientTape() as tape:
        last_conv_layer_output, preds = grad_model(img_array)
        # 우리의 모델은 num_classes개의 예측값을 출력한다. 각 예측값은 데이터가 해당 클래스 확률값이다.
        # 미분 (더 정확히는 그레디언트)는 한 개의 예측값에 대해서 계산할 수 있기 때문에
        # 가장 확률값이 큰 예측값에 대응하는 클래스 (혹은 함수의 인자로 전달한 클래스)만을 고려할 것이다. -> class_channel
        if pred_index is None:
            pred_index = tf.argmax(preds[0])
        class_channel = preds[:, pred_index]

    # 이 친구가 바로 미분 계산이다. 
    # 모델의 어떤 한 클래스에 대한 예측값(class_channel) 을 마지막 feature maps(last_conv_layer_output)으로 미분한다.
    grads = tape.gradient(class_channel, last_conv_layer_output)

    # 마지막 feature maps은 여러 개의 채널을 갖고 있다. (현재 [1, 높이, 너비, 채널] 상태, 1은 데이터 개수)
    # 각 채널마다 등장한 미분값들을 평균 내려주자. (채널의 중요도를 결정)
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))

    # 우리는 마지막 feature maps에서 중요한 위치를 알고 싶다. 
    # 하지만 채널 개수가 많기 때문에 어떤 위치가 중요한지 결정해줘야 한다.
    # 채널의 중요도를 가중치로 가중합을 해주자 ([높이, 너비, 1]가 됨)
    last_conv_layer_output = last_conv_layer_output[0]
    # @는 행렬곱 연산이다. 축을 잘 맞춰줌으로써 행렬곱이 곧 가중합 연산이 되도록 한 것이다.
    heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis] 
    heatmap = tf.squeeze(heatmap)  # 축 쥐어짜서 (squeeze) 없애주기 ([높이, 너비, 1] -> [높이, 너비]가 됨)

    # 0과 1사이 값으로 만들어주기
    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
    return heatmap.numpy()

heatmap = make_gradcam_heatmap(X_test[0:1], model, 'max_pooling2d_1')
plt.matshow(heatmap)
plt.show()

def display_gradcam_with_img(img, heatmap, alpha=0.4):
    # 0과 1 사이의 값을 갖는 행렬 -> 0과 255 사이의 정수를 갖는 행렬
    heatmap = np.uint8(255 * heatmap)

    # 이쁜 히트맵을 위한 칼라맵 설정
    jet = cm.get_cmap("jet")
    jet_colors = jet(np.arange(256))[:, :3]  # 0~255에 대응하는 RGB 값 불러오기 [256, 3]
    jet_heatmap = jet_colors[heatmap]  # RGB 이미지가 됨 (왜 그렇게 되는지는 굉장히 어려워 보임 ...)

    jet_heatmap = keras.preprocessing.image.array_to_img(jet_heatmap)  # 행렬을 이미지로 바꿔주기 (이미지 크기 바꾸기 위해)
    jet_heatmap = jet_heatmap.resize((img.shape[1], img.shape[0]))     # 히트맵 사이즈를 이미지 크기로 바꿔주기
    jet_heatmap = keras.preprocessing.image.img_to_array(jet_heatmap)  # 이미지를 행렬로 바꿔주기

    superimposed_img = jet_heatmap * alpha + img  # 히트맵과 이미지 겹치기 (덧셈 연산으로 가능)
    superimposed_img = np.uint8(superimposed_img)

    # 그림 그려
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    axes[0].imshow(img)  # 원본 이미지 그려
    axes[1].imshow(heatmap)  # 히트맵 그려
    axes[2].imshow(superimposed_img)  # 겹친 이미지 그려
    for ax in axes:
        ax.axis('off')  # 축에 눈금이랑 숫자 다 없애

img = X_test[1]
heatmap =  make_gradcam_heatmap(img[np.newaxis], model, 'max_pooling2d_1')
k = display_gradcam_with_img(255 * img, heatmap, alpha=0.8)