※ 본 게시글은 필자의 개인 경험을 토대로 작성된 것으로, 잘못된 정보를 포함하고 있을 수 있습니다. ※
필자의 경우 ImageNet 데이터 세트(대표적인 이미지 데이터 세트)을 활용한 실험을 진행하고 있다. 그래서, 특정한 ImageNet 학습 데이터가 주어졌을 때, 해당 이미지에 대하여 CAM (Class Activation Map) 이미지을 뽑은 뒤에, 해당 activated 영역에 대하여 이미지 processing을 수행하여, 변경된 이미지를 얻는 방법에 대하여 다룬다.
우리는 CAM을 사용하면, 모델이 특정한 결과(output)을 내뱉을 때, 이미지에서 어떤 영역(region)을 보고 해당 결과를 내뱉었는지 해당 activation map을 얻을 수 있다. 즉, 특정한 클래스(class)로 분류를 했을 때, 해당 클래스로 분류를 하게 된 이유를 확인할 수 있다. 가장 먼저, 다음과 같은 코드로 이미지넷(ImageNet) 1,000 classes 데이터 세트를 불러올 수 있다.
※ 데이터 세트 불러오기 ※
ImageNet-1k 데이터 세트를 불러올 수 있다.
from datasets import load_dataset
imagenet_dataset = load_dataset("imagenet-1k")
해당 데이터 세트는 Hugging Face에서 제공하고 있기 때문에, PyTorch의 데이터 로더로 불러오기 위해서는 다음과 같이 CustomDataset과 같은 클래스를 별도로 작성할 필요가 있다.
import torch
class CustomDataset(torch.utils.data.Dataset):
def __init__(self, train_mode=True, transforms=None):
if train_mode:
self.dataset = imagenet_dataset["train"]
else:
self.dataset = imagenet_dataset["validation"]
self.transforms = None
if transforms:
self.transforms = transforms
def __getitem__(self, index):
image = self.dataset[index]["image"]
label = self.dataset[index]["label"]
current = image.convert("RGB")
if self.transforms:
current = self.transforms(current)
return current, label
def __len__(self):
return len(self.dataset)
이후에 다음과 같이 기본적인 전처리(pre-processing) 소스 코드를 작성할 필요가 있다. 일반적으로 ImageNet에 대하여 데이터를 단순히 불러와 모델에 넣을 때는 이미지 크기를 (224 X 224) 크기로 변경하고, 정규화(normalization)를 수행하여 최종적으로 처리된 이미지를 모델에 넣어 학습을 진행할 수 있다.
import torchvision.transforms as transforms
from pytorch_ood.utils import ToRGB
mean = [x for x in [0.485, 0.456, 0.406]]
std = [x for x in [0.229, 0.224, 0.225]]
trans = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
ToRGB(),
transforms.ToTensor(),
transforms.Normalize(std=std, mean=mean)
])
필자는 미리 준비된 ImageNet의 1,000개의 클래스마다 각 100개씩 이미지를 가지고 있는 간단한 형태의 데이터 세트를 준비하여, 거기에서 불러오도록 했다. 참고로 test_dataset은 실제로 사용되지 않지만, 한 번 불러와서 확인해 보았다.
from torchvision.datasets import ImageFolder
batch_size = 1
train_dataset = ImageFolder("./imagenet_100_images_per_class", transform=trans)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=16)
test_dataset = CustomDataset(train_mode=False, transforms=trans)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=16)
print("Size of training dataset:", len(train_dataset))
print("Size of test dataset:", len(test_dataset))
※ 딥러닝 모델 불러오기 및 대상 레이어 설정하기 ※
이후에 필자는 LayerCAM을 사용하기로 하였고, 다음과 같이 ResNet50 모델에서 layer2을 대상으로 동작하도록 설정을 진행했다. 참고로 입력 데이터의 크기 및 대상 레이어(target layer)를 파라미터로 넣도록 설정되어 있다.
from torchcam.methods import LayerCAM
num_classes = 1000
img_size = 224
input_shape = (3, img_size, img_size)
model = resnet50(num_classes=num_classes, pretrained=True).cuda().eval()
target_layer = model.layer2
localize_net = LayerCAM(model, target_layer=target_layer, input_shape=input_shape)
※ CAM을 이용해 이미지 처리하기 ※
필자는 다음과 같이 존재하는 모든 이미지에 대하여 activation map을 얻은 뒤에, activation 값이 0.2보다 큰 영역(중요한 영역들)을 이미지에서 제거한 뒤에 in-painting 방법을 이용하여 해당 내용을 채우는 방식으로 일종의 의미 없는(특정 클래스에 대한 정보가 없는) 이미지 데이터를 만드는 작업을 해보았다.
import cv2
from tqdm import tqdm
import numpy as np
cam_lambda = 0.2
save_dir = os.path.join(f"./saved_samples")
for class_idx in range(num_classes):
os.makedirs(os.path.join(save_dir, str(class_idx)), exist_ok=True)
for batch_idx, (data, target) in enumerate(tqdm(train_loader)):
if torch.cuda.is_available():
data = data.cuda()
target = target.cuda()
# 이미지를 모델에 넣어 결과 얻기
out = model(data)
# 모델이 출력한 결과를 토대로 Class Activation Map (CAM) 얻기
activation_map = localize_net(out.squeeze(0).argmax().item(), out)
activation_map = activation_map[0].squeeze().detach().cpu().numpy()
# activation map의 해상도(resolution)을 원래 이미지와 동일하게 만들기
if activation_map.shape[0] != img_size:
x = cv2.resize(activation_map, (img_size, img_size))
else:
x = activation_map
activation_map = x
# 원본 이미지를 NumPy 객체로 저장하기
x_data_array = np.transpose(data.detach().cpu().numpy(), [0, 2, 3, 1])
origin_x_data = (x_data_array * np.array(std).reshape([1, 1, 1, 3])) + np.array(mean).reshape([1, 1, 1, 3])
origin_x_data = np.uint8(origin_x_data * 255)[0]
# 결과적으로 얻은 activation map이 특정한 threshold보다 낮은 부분을 마스크 이미지로 얻기
background_mask = np.uint8(activation_map < cam_lambda)
# 원본 이미지에 대하여 마스크 영역을 지운 이미지 얻기
remove_image = np.copy(origin_x_data) * np.expand_dims(background_mask, axis=-1)
# inpainting을 진행하기 위하여 마스크 이미지 생성
target_mask = -1 * (background_mask.astype(np.float32) - 1.)
# FS (Fast Marching) inpainting 방식을 적용하여 이미지 생성
inpaint = cv2.inpaint(remove_image, target_mask.astype(np.uint8), 5, cv2.INPAINT_TELEA)
# 최종적인 결과를 이미지 파일 형태로 저장하기
class_idx = target.detach().cpu().numpy().flatten()[0]
# ImageFolder 라이브러리를 사용하는 경우 클래스(class) 번호가 매칭 이슈로, class_to_idx가 필요할 수 있음
# 이때, class_to_idx는 폴더 이름(클래스 이름)으로 인덱스(index)를 찾는 dictionary 객체를 의미
save_original_train_path = os.path.join(save_dir, str(class_idx), f"{batch_idx}_{class_idx}_{str(cam_lambda)}_original.png")
save_ood_train_path = os.path.join(save_dir, str(class_idx), f"{batch_idx}_{class_idx}_{str(cam_lambda)}_train.png")
save_ood_mask_path = os.path.join(save_dir, str(class_idx), f"{batch_idx}_{class_idx}_{str(cam_lambda)}_mask.png")
cv2.imwrite(save_original_train_path, origin_x_data[..., ::-1].astype(np.uint8))
cv2.imwrite(save_ood_mask_path, (target_mask * 255).astype(np.uint8))
cv2.imwrite(save_ood_train_path, inpaint[..., ::-1].astype(np.uint8))
예를 들어 코끼리(elephant) 클래스에 속하는 이미지가 주어졌을 때, 다음과 같이 CAM을 사용하여 코끼리의 뿔과 같은 위치를 정확히 잘 잡아내는 것을 확인할 수 있다. 이후에 inpaint() 함수를 이용하여 코끼리의 특성을 지울 수 있다.
결과적으로 총 10만개의 ImageNet 데이터 세트에 대하여 CAM을 적용할 수 있었다. 10만 장에 대하여 총 2시간 30분 정도 소요되었다. 이러한 과정에서 하나의 TITAN RTX GPU만을 사용했으며, 배치 크기(batch size)를 1로 사용하여 2GB 이하의 VRAM을 사용하여 간단하게 돌릴 수 있었다. 이후에 다음의 코드를 사용하여 결과를 확인할 수 있다.
import os
# 각 클래스(class) 레이블 정보를 하나씩 확인하며
for i in range(1000):
directory = f"{save_dir}/{i}"
cnt = 0
# 폴더 내부의 경로를 하나씩 확인하며
for path in os.listdir(directory):
# 현재의 경로(path)가 파일인 경우 카운트
if os.path.isfile(os.path.join(directory, path)):
cnt += 1
# 각 레이블에 포함된 파일의 개수 출력
print(f"The number of file: {cnt} for the label {i}")
프로그램 실행 결과 예시는 다음과 같다. 각 이미지에 대하여 (1) original.png, (2) train.png, (3) mask.png로 3개의 이미지가 결과적으로 저장되는 것을 확인할 수 있다.
