8. Learning Deep Features for Discriminative Localization(CAM) - code review

 

안녕하세요.

 

 

오늘은 저번 글에서 review 했던 Learning Deep Features for Discriminative Localization 논문의 code review를 진행해보겠습니다.

 

 

해당 논문 실험은, 기존에 제가 실험했던 CIFAR10이나 MNIST와 같이 torchvision에서 제공하는 데이터셋이 아닌 Kaggle에서 제공하는 CatVsDog dataset을 사용하여서 dataset을 setting 하는 작업이 별도로 필요합니다.

 

 

따라서 실제로 해당 코드를 직접 구동해보시면서 따라가실 분들은 제 Github에 있는 코드를 보시고 .ipynb 파일과 README 내용들을 따라가시면서 직접 dataset을 setting 하신 후 실험을 진행해주시면 되겠습니다.

 

그리고 Class Activation Map은 모델과, 이미 학습이 완료된 weight가 있어야 만들 수 있습니다.

 

따라서 실제 생성을 하시기 위해서는 모델을 별도로 학습해주셔야 합니다.

 

이 과정 또한 제 Github에 나와있으니, 실제 구현을 해보실 것이라면 이 과정도 거치셔야 합니다.

 

제 Github 주소는 다음과 같습니다.

 

github.com/PeterKim1/paper_code_review/tree/master/8.%20Learning%20Deep%20Features%20for%20Discriminative%20Localization(CAM)

PeterKim1/paper_code_review

 

dataset이 setting 되어있고, 이미 학습이 완료된 weight가 있다는 가정 하에 모델과 관련한 내용만 review를 진행해보겠습니다.

 

 

그럼 시작하겠습니다.

 

 

 

1. 저장된 모델 불러오기

 

 

 

먼저, Class Activation Map은 이미 학습이 완료된 model을 기반으로 만들어지기 때문에 이미 학습이 완료된 모델이 있어야 하며 이를 불러와야 만들 수 있습니다.

 

이를 불러오는 과정을 진행합니다.

 

import torch
import os

sample = torch.load(os.path.join(os.getcwd(), '2021-05-04-13:55/checkpoint/ckpt.pth'))

import torchvision.models as models
import torch.nn as nn

resnet18 = models.resnet18(pretrained = False)

resnet18.fc = nn.Linear(512, 2)
print(resnet18.fc)

resnet18.load_state_dict(sample['net'])

 

먼저 저장된. pth 파일을 sample이라는 변수로 저장합니다.

 

그리고 torchvision에서 지원하는 resnet18 모델을 가져옵니다. 

 

이때, torchvision의 resnet18의 경우 1000 class classification 모델이므로 마지막 layer를 binary classification layer로 변경합니다.

 

그러고 나서 load_state_dict()를 사용해 저장된 weight를 가져옵니다.

 

load_state_dict()를 사용하고 나서 다음과 같은 메시지가 떠야 정상적으로 불러오기가 된 것입니다.

 

 

 

 

2. DataLoader setting 및 모델을 GPU에 올리기

 

 

Class Activation Map은 내가 어떤 이미지를 분류하는 과정에서, 이미지의 어떤 discriminative region을 보고 판단했는지를 보여주는 방법론입니다.

 

따라서, 기본적으로는 처음 보는 이미지를 classification 합니다. 

 

그래서 train dataloader는 필요가 없으며, test dataloader만 가져오면 됩니다.

 

import cv2
import torchvision
import torchvision.transforms as transforms
import datetime
import numpy as np

device = 'cuda' if torch.cuda.is_available() else 'cpu'

transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.4876, 0.4544, 0.4165), (0.2257, 0.2209, 0.2212)),
])

test_dataset = torchvision.datasets.ImageFolder(root = os.path.join(os.getcwd(), 'CatvsDog_test'), transform = transform_test)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = 1, shuffle = True, num_workers = 0)


current_time = datetime.datetime.now() + datetime.timedelta(hours= 9)
current_time = current_time.strftime('%Y-%m-%d-%H:%M')


saved_loc = os.path.join('./', current_time)
if os.path.exists(saved_loc):
    shutil.rmtree(saved_loc)
os.mkdir(saved_loc)

print("결과 저장 위치: ", saved_loc)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

resnet18.to(device)

 

test dataset을 가지고 torchvision.datasets.ImageFolder로 dataset 객체를 만들어준 뒤에, 이를 DataLoader에 투입시킵니다.

 

이때, batch_size는 1로 주어 한 개의 이미지씩 나오도록 만들어줍니다.

 

물론 여러 장씩 나오게 할 수도 있지만, 1개씩 나오게끔 하는 것이 전반적으로 간단해질 것 같아 이렇게 설정하였습니다.

 

그리고 결과를 저장해주기 위해서, saved_loc이라는 이름으로 현재 실행 시간을 폴더명으로 해서 폴더를 만들어줍니다.

 

마지막으로, 불러온 resnet18 모델을 GPU로 이동시킵니다.

 

 

3. CAM 만들기

 

 

본격적으로 CAM을 만들어보겠습니다.

 

import torch.nn.functional as F
import matplotlib.pyplot as plt
from scipy.ndimage import label
from skimage.measure import regionprops
import matplotlib

# final conv layer name 
finalconv_name = 'layer4'

# inference mode
resnet18.eval()

# number of result
num_result = 3

먼저 CAM을 만드는데 필요한 library를 가지고 오고, 마지막 conv layer의 이름을 가져옵니다.

 

그리고 resnet18을 inference mode로 변경한 뒤, result를 몇 개 나오게 만들 것인지를 결정합니다.

 

feature_blobs = []

# output으로 나오는 feature를 feature_blobs에 append하도록
def hook_feature(module, input, output):
    feature_blobs.append(output.cpu().data.numpy())
    
resnet18._modules.get(finalconv_name).register_forward_hook(hook_feature)

# get the softmax weight
params = list(resnet18.parameters())
weight_softmax = np.squeeze(params[-2].cpu().detach().numpy()) # [2, 512]

 

feature_blobs는 빈 형태의 list인데, 추후 여기에 feature가 들어가게 됩니다.

 

hook_feature라는 함수는 output으로 나오는 feature를 feature_blobs에 append 하도록 짜인 함수입니다.

 

그리고 resnet18._modules.get(finalconv_name).register_forward_hook(hook_feature)라는 코드를 통해서 마지막 conv layer인 'layer4'를 거쳐서 나오게 되는 feature를 feature_blobs에 저장하도록 만들어줍니다.

 

resnet18의 parameter들을 list로 만들어 params에 저장해주고, 여기서 마지막에서 두 번째 weight를 weight_softmax라는 변수로 저장해줍니다.

 

resnet18의 parameter는 layer 순서대로 저장이 되어 있는데, 가장 마지막은 fully-connected layer의 bias로 구성되어 있고 마지막에서 두 번째는 fully-connected layer의 가중치가 저장되어 있습니다.

 

따라서 params [-2]를 weight_softmax에 저장해주게 됩니다.

 

만약 제가 사용한 모델이 아니라 별도의 모델을 사용하시는 경우, 본인의 모델의 parameter들 중에서 어떤 부분이 output을 내는데 이용되는 parameter인지를 확인하신 후 사용하셔야 합니다.

 

 

# generate the class activation maps
def returnCAM(feature_conv, weight_softmax, class_idx):
    size_upsample = (224, 224)
    _, nc, h, w = feature_conv.shape # nc : number of channel, h: height, w: width (1, 512, 7, 7)
    output_cam = []
    # weight 중에서 class index에 해당하는 것만 뽑은 다음, 이를 conv feature와 곱연산
    cam = weight_softmax[class_idx].dot(feature_conv.reshape((nc, h*w))) 
    cam = cam.reshape(h, w)
    cam = cam - np.min(cam)
    cam_img = cam / np.max(cam)
    cam_img = np.uint8(255 * cam_img)
    output_cam.append(cv2.resize(cam_img, size_upsample))
    return output_cam

 

먼저 CAM을 만드는 데 사용되는 함수부터 소개해드리겠습니다.

 

returnCAM이라는 함수는 마지막 conv layer를 통과해서 나오게 되는 feature와 마지막 output을 만드는 데 사용되는 weight, 그리고 해당 예측의 class index를 input으로 받아 CAM을 만들어주는 함수입니다.

 

먼저 어떤 사이즈로 CAM을 만들 것인지를 결정합니다. 

 

저희 이미지는 (224, 224)로 구성되어 있어 이를 이미지의 사이즈와 동일하게 맞춰주면 됩니다.

 

다음으로는 feature_conv.shape를 이용해 channel의 수와 height, width를 저장합니다.

 

output_cam은 빈 list로 추후 결과를 저장하는 list이며

 

cam은 weight_softmax [class_idx]와 feature_conv를 (nc, h*w)로 reshape 한 것을 곱 연산해줍니다.

 

weight_softmax의 경우 (2, 512)로 구성되는데요, 이 중에서 class index에 해당하는 값만 선택합니다.

 

그렇다면 shape가 (1, 512)가 되겠죠? 

 

그러고 feature_conv는 (1, 512, 7, 7)의 shape를 가지고 있는데 이를 (512, 7*7) 형태로 바꿔주고 나서 dot product를 수행합니다.

 

matrix 연산은 첫 번째 matrix의 가로와 두 번째 matrix의 세로가 같아야만 연산이 되니까요. 그래서 결과는 (1, 7*7)이 됩니다.

 

이 결과를 (7, 7) 모양으로 다시 reshape 해주면 마지막 conv layer를 통과해서 나온 (1, 512, 7, 7)와 가로, 세로 관점에서 동일한 크기를 가지는 CAM을 만들게 되는 것이죠.

 

그러고 나서 이를 최솟값으로 빼주고, 최댓값으로 나눠줍니다. 이 작업은 보통 min-max scaling이라고 부르던 것 같은데 이 작업을 수행하게 되면 가장 작은 값은 0이 되고 가장 큰 값은 1이 될 겁니다.

 

그러고 나서 여기에 255를 곱해줘서 사실상 0부터 255의 값을 가지는 image를 만들어낸 것이라고 보시면 됩니다.

 

마지막으로는 이를 cv2.resize()를 활용해 이미지의 원래 사이즈인 (224, 224) 형태로 resize 해줍니다.

 

def generate_bbox(img, cam, threshold):
    labeled, nr_objects = label(cam > threshold)
    props = regionprops(labeled)
    return props

 

해당 함수는 CAM을 기반으로 bounding box를 그리는 함수인데, 자세한 내용은 해당 글을 참고해주시면 됩니다.

 

이미 제가 글을 올렸던 내용이라서 자세한 내용은 글로 대체합니다.

 

cumulu-s.tistory.com/38

scipy.ndimage.label과 skimage.measure.regionprops를 이용한 이미지 내에서 특정값 이상의 값을 가지는 위치

 

for i, (image, target) in enumerate(test_loader):
    
    # 모델의 input으로 주기 위한 image는 따로 설정
    image_for_model = image.clone().detach()

    # Image denormalization, using mean and std that i was used.
    image[0][0] *= 0.2257
    image[0][1] *= 0.2209
    image[0][2] *= 0.2212
    
    image[0][0] += 0.4876
    image[0][1] += 0.4544
    image[0][2] += 0.4165

 

이제부터 본격적으로 CAM 만들기에 들어갑니다.

 

먼저 test_loader로부터 image와 target을 받고, 모델의 input으로 주기 위한 image를 따로 image_for_model로 지정합니다.

 

그리고 이전에 학습을 진행할 때와 test loader에도 동일하게 normalization을 해줬어서 normalized 되지 않은 이미지를 저장해주기 위해 이를 반대로 역정규화 해줍니다. 

 

이렇게 하지 않으면 원본 이미지를 저장할 때 굉장히 이상한 이미지의 모습으로 보이기 때문에 정상적인 이미지 형태를 저장하고자 이 작업을 추가하였습니다.

 

    fig = plt.figure(tight_layout = True)
    rows = 1
    cols = 3
    
    ax1 = fig.add_subplot(rows, cols, 1)
    ax1.imshow(np.transpose(image[0].numpy(), (1, 2, 0)))
    ax1.set_title("Original Image")
    ax1.axis("off")

 

원래 이미지와 작업을 수행한 결과를 저장해야 하므로, plt.figure로 fig를 만들어주고 row는 1, col은 3으로 잡습니다.

 

그리고 여기에 add_subplot을 통해서 원본 이미지를 저장해줍니다. np.transpose의 경우 차원의 위치를 바꿔주는데, pytorch에서 지원하는 차원의 순서와 matplotlib에서 지원하는 차원의 순서가 다르기 때문에 이 변환 과정을 거쳐야 합니다.

 

    image_PIL = transforms.ToPILImage()(image[0])
    image_PIL.save(os.path.join(saved_loc, 'img%d.png' % (i + 1)))

 

해당 코드는 image를 PIL Image로 읽은 다음, 저장하는 코드입니다.

 

    # 모델의 input으로 사용하도록.
    image_tensor = image_for_model.to(device)
    logit = resnet18(image_tensor)
    h_x = F.softmax(logit, dim=1).data.squeeze()
    
    probs, idx = h_x.sort(0, True)
    print("True label : %d, Predicted label : %d, Probability : %.2f" % (target.item(), idx[0].item(), probs[0].item()))

 

아까 저장해둔 image_for_model을 GPU로 이동하고, 이를 resnet18에 통과시켜서 결과를 내놓습니다. 

 

해당 이미지가 개와 고양이 중 무엇인지를 예측한 결과가 logit에 저장됩니다.

 

이에 대해서 softmax를 거쳐주고, 이것을 sorting 한 결과를 probs, idx에 저장합니다.

 

probs는 예측 확률이 될 것이고, idx는 예측이라고 판단된 결과의 index가 되겠죠. 

 

이를 이용해서 실제 label과 예측된 label, 그리고 probability를 출력할 수 있게 됩니다.

 

    CAMs = returnCAM(feature_blobs[0], weight_softmax, [idx[0].item()])
    img = cv2.imread(os.path.join(saved_loc, 'img%d.png' % (i + 1)))
    height, width, _ = img.shape
    heatmap = cv2.applyColorMap(cv2.resize(CAMs[0], (width, height)), cv2.COLORMAP_JET)
    result = heatmap * 0.7 + img * 0.5
    cv2.imwrite(os.path.join(saved_loc, 'cam%d.png' % (i + 1)), result)

 

아까 feature_blobs는 마지막 conv layer를 통과하고 나서 나오게 되는 feature라는 말씀을 드렸었습니다.

 

이는 (1, 512, 7, 7)의 크기를 가지고 있는데, 이를 returnCAM이라는 함수에 투입시켜 Class Activation Map을 만들어줍니다.

 

그리고 img라는 변수명으로 원본 이미지를 읽어오고, heatmap이라는 이름으로 Class Activation Map을 저장합니다. 

 

result는 최종 결과이며, CAM와 image의 가중치의 합으로 계산합니다.

 

heatmap의 가중치가 클수록 사진에서 CAM이 더 뚜렷하게 나타나고, img의 가중치가 클수록 사진에서 원본 사진이 더 뚜렷하게 나타납니다.

 

그리고 이 최종 결과인 result를 저장합니다.

 

    props = generate_bbox(img, CAMs[0], CAMs[0].max()*0.6)
    
    ax2 = fig.add_subplot(rows, cols, 2)
    im2 = ax2.imshow(CAMs[0])
    ax2.set_title("Class Activation Map")
    ax2.axis("off")
    
    ax3 = fig.add_subplot(rows, cols, 3)
    ax3.imshow(np.transpose(image[0].numpy(), (1, 2, 0)))
    ax3.set_title("Original Image + patch")
    ax3.axis("off")

 

generate_bbox 함수를 이용해서 props를 만들어줍니다.

 

CAM은 ax2라는 이름으로 subplot을 만들어주고, 원래 이미지에 patch를 추가한 이미지를 ax3라는 이름으로 만들어줍니다.

 

ax3에는 아직 patch를 추가하진 않았고, 현재는 original image만 올라가 있는 상태입니다.

 

    for b in props:
        bbox = b.bbox
        xs = bbox[1]
        ys = bbox[0]
        w = bbox[3] - bbox[1]
        h = bbox[2] - bbox[0]
        
        rect = matplotlib.patches.Rectangle((xs, ys), w, h, linewidth = 2, edgecolor = 'r', facecolor = 'none')
        ax3.add_patch(rect)
    
    
    plt.savefig(os.path.join(saved_loc, 'Result%d.png' % (i+1)))
    
    plt.show()
    
    if i + 1 == num_result:
        break
        
    feature_blobs.clear()

 

for b in props:를 통해서 ax3에 bounding box를 만들어줍니다.

 

해당 코드의 자세한 내용은 위에 올렸던 제 블로그 링크를 확인해주시면 될 것 같습니다.

 

만들어진 plt를 save 하고, show() 합니다.

 

만약 i + 1가 num_result랑 같아지면, 해당 for문을 break 하고요.

 

feature_blobs는 for문을 돌릴 때마다 비워줍니다.

 

이를 통해서 나오게 되는 결과는 다음과 같습니다.

 

 

 

Class Activation Map을 보면 고양이가 위치한 부분과 개가 위치한 부분에서 더 많이 activation 되는 것을 확인할 수 있으며, 이를 통해서 만들어진 bounding box도 각각 개와 고양이에 해당하는 부분만 잡힌다는 사실을 확인할 수 있습니다.

 

 

 

여기까지 Class Activation Map에 대한 code review를 진행해보았습니다.

 

실제로 구현해보니 정말 각 동물에 해당하는 이미지의 위치에 활성화된다는 것을 확인할 수 있었습니다.

 

이번 review는 여기까지 해서 마무리짓고, 다음 글에서 뵙도록 하겠습니다.

 

감사합니다.