[Deep Learning] DETR 모델 이해하고 실습하기 (2)

안녕하세요!

지난 실습에서는 pytorch로 DETR 모델을 구현해보았습니다. 이번에는 이 모델에 제가 찍은 사진을 입력해서 객체 탐지를 해보겠습니다. 구현한 모델은 지난 실습 포스팅을 참고해주세요!

https://jun-eon.tistory.com/entry/%EB%94%A5%EB%9F%AC%EB%8B%9DDETR-%EB%AA%A8%EB%8D%B8-%EC%9D%B4%ED%95%B4%ED%95%98%EA%B3%A0-%EC%8B%A4%EC%8A%B5%ED%95%98%EA%B8%B0

[Deep Learning] DETR 모델 이해하고 실습하기 (1)

 

import math

from PIL import Image # 이미지 파일을 열거나 회전시키거나 편집하는 등의 작업에 활용
import matplotlib.pyplot as plt # 데이터 시각화에 활용
%config InlineBackend.figure_format = 'retina' # 이미지 품질을 높여 고해상도 이미지를 얻는 데 활용

import ipywidgets as widgets
from IPython.display import display, clear_output # 셀 출력 결과를 간략하게 표시

import torch
from torch import nn


from torchvision.models import resnet50
import torchvision.transforms as T	# 이미지 전처리 작업
torch.set_grad_enabled(False);

 

먼저 실습에 필요한 라이브러리를 불러와줍니다. 

# COCO classes
CLASSES = [
    'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A',
    'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
    'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack',
    'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
    'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
    'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
    'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A',
    'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
    'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A',
    'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
    'toothbrush'

데이터는 COCO  데이터셋을 활용했습니다. COCO 데이터셋은 다양한 객체 종류를 포함하는 대규모 이미지 데이터셋입니다. COCO 데이터셋의 일부 인덱스는 특정 클래스와 매핑되지 않기 때문에 'N/A'로 표시됩니다. 

# colors for visualization
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
          
# standard PyTorch mean-std input image normalization
transform = T.Compose([
    T.Resize(800),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

COLORS는 객체 탐지에서 바운딩 박스를 표시할 때 사용할 색상을 정의한 것입니다. 

transform 부분은 이미지 전처리를 위한 부분인데요, T.Resize(800)을 통해 이미지의 가장 긴 변이 800픽셀이 되도록 조정하고, T.ToTensor()를 통해 이미지를 텐서 형태로 변환합니다. T.Normalize 부분은 이미지의 각 채널(R, G, B)에 대해 평균과 표준편차로 정규화하는 부분입니다. 

# for output bounding box post-processing

def box_cxcywh_to_xyxy(x):
    x_c, y_c, w, h = x.unbind(1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
         (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=1)

def rescale_bboxes(out_bbox, size):
    img_w, img_h = size
    b = box_cxcywh_to_xyxy(out_bbox)
    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
    return b

def plot_results(pil_img, prob, boxes):
    plt.figure(figsize=(16,10))
    plt.imshow(pil_img)
    ax = plt.gca()
    colors = COLORS * 100
    for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):
        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
                                   fill=False, color=c, linewidth=3))
        cl = p.argmax()
        text = f'{CLASSES[cl]}: {p[cl]:0.2f}'
        ax.text(xmin, ymin, text, fontsize=15,
                bbox=dict(facecolor='yellow', alpha=0.5))
    plt.axis('off')
    plt.show()

첫 번째 함수는 바운딩 박스의 형식을 (cx, cy, w, h) 형식에서 (xmin, ymin, xmax, ymax) 형태로 변환하는 함수입니다. 

두 번째 함수는 모델이 예측한 바운드 박스를 실제 이미지의 크기에 맞게 조정하는 함수입니다. 

세 번째 함수는 모델의 예측 결과를 시각화하여 보여주는 함수입니다. 

model = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=True)
model.eval();

DETR 모델을 외부에서 불러오는 코드입니다. model.eval() 을 통해 모델을 평가 모드로 설정합니다. 학습 시에만 사용되는 dropout이나 batch normalization과 같은 동작을 비활성화합니다. 

 

다음으로는 제 이미지를 업로드하고 모델에 입력해보겠습니다. 

from google.colab import drive

drive.mount('/content/drive')

im = Image.open('/content/drive/MyDrive/24-2/딥러닝/images/my_image_1.jpeg') 
im = im.rotate(-90, expand=True)

저는 코랩으로 작업을 하고 있어서 구글 드라이브에 이미지를 업로드해서 활용했습니다. 편한 방법을 사용하셔서 im 변수에 이미지를 할당해주세요. 저는 제가 여행갔을 때 찍었던 사진을 활용해보았습니다. 탐지할 객체가 골고루 있는 것 같아서 이 사진으로 선정했습니다. 

아이슬란드 어딘가

# mean-std normalize the input image (batch-size: 1)
# 기존에 해놓은 설정대로 입력 이미지에 대한 전처리를 수행. 
img = transform(im).unsqueeze(0)

# propagate through the model
# 모델을 통한 예측
outputs = model(img)

# keep only predictions with 0.7+ confidence
# 예측 값을 softmax를 통해 클래스별 확률로 변환
probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]
# 객체에 대해 가장 높은 확률을 가진 클래스가 0.9 이상인 경우만 필터링하여 신뢰도가 높은 객체만 출력
keep = probas.max(-1).values > 0.9

# convert boxes from [0; 1] to image scales
# bounding box를 객체의 크기에 맞게 조정
bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)

# 결과 시각화
plot_results(im, probas[keep], bboxes_scaled)

다음과 같은 결과를 얻을 수 있었습니다. 

 

매우 신기하다!

다음 실습에서는 해당 모델이 특정 클래스와 바운딩 박스를 예측할 때, 이미지의 어떤 부분에 집중(attention)했는지 알아보겠습니다!

참조:

https://github.com/facebookresearch/detr

GitHub - facebookresearch/detr: End-to-End Object Detection with Transformers