R-CNN (2014) PyTorch Implementation

R-CNN을 Pytorch를 활용하여 구현하고자 한다.

https://arsetstudium.tistory.com/33에서 공부한 내용을 토대로 구현보면 아래와 같다.

 이때, R-CNN의 CNN 구조는 AlexNet과 같으므로 생략하고 기존의 classification과 다른 부분에 집중하고자 한다.

 

바로 bounding box의 생성 부분인데 우선 ROI (Region of Interest)라는 개념도 알아야 한다.

 

ROI는 이름 그대로 어떤 이미지에서 관심이 가는 특정 대상에 대한 영역을 의미한다.

Object detection에서 ROI 지정은 bounding box로 하며, 아래의 Figure 1를 예시로 알아보자.

 

Figure 1. Bird and ROI box

 

우선 interest는 새다. 그리고 그 새에 대한 영역인 ROI를 빨간색 bounding box (이하 bbox)로 표시했음을 알 수 있다.

 

솔직하게 말해서 해당 부문에 대한 코드를 직접 작성할 수 있다는 자신이 없으므로 해당 포스트에서는 bounding box를 만드는 PyTorch 코드를 이용해서 최대한 라인 바이 라인으로 분석하며 공부하고자 한다.

 

Drawing Bounding Box Example Code

아래의 링크에 나온 코드를 분석하고자 한다.

https://pytorch.org/vision/main/auto_examples/others/plot_repurposing_annotations.html#sphx-glr-auto-examples-others-plot-repurposing-annotations-py

 

맨 처음에는 필요한 라이브러리를 임포트하고 bounding box와 segmentation을 그리기 위한 show 함수를 선언하게 된다.

import os
import numpy as np
import torch
import matplotlib.pyplot as plt

import torchvision.transforms.functional as F


plt.rcParams["savefig.bbox"] = "tight"

def show(imgs):
    if not isinstance(imgs, list):
        imgs = [imgs]
    fix, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    for i, img in enumerate(imgs):
        img = img.detach()
        img = F.to_pil_image(img)
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

 

 

그 다음 object detection에 필요한 bounding box (이하 bbox)를 그리기 전에 image와 mask를 불러온다.

이때 mask는 (num_objects, height, width)로 구성되는데, num_objects는 이미지에 있는 object의 개수,

height와 width는 이미지의 크기다.

 

가령 (4, 224, 224)의 정보가 있다면 224 x 224 이미지에 4개의 object가 있다는 뜻이다.

 

from torchvision.io import read_image

img_path = os.path.join(data_path, 'PNGImages', "FudanPed00054.png")
mask_path = os.path.join(data_path, 'PedMasks', "FudanPed00054_mask.png")

# 이미지 로드
img = read_image(img_path)
# mask 로드
mask = read_image(mask_path)

# 출력으로 이미지와 마스크 확인
print(mask.size())
print(img.size())
print(mask)

>>
torch.Size([1, 498, 533]) # 498 x 533 image에 1개의 object가 있다.
torch.Size([3, 498, 533]) # 498 x 533 사이즈의 RGB image다.
tensor([[[0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         ...,
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0]]], dtype=torch.uint8)

 

 

이제 이미지를 matplotlib를 통해서 출력하여 확인한다.

# 불러온 이미지를 출력해서 확인한다.
import matplotlib.pyplot as plt
plt.imshow(F.to_pil_image(img))

 

Figure 1. FundanPed00054

 

위의 이미지가 출력되어야 한다. 해당 이미지에는 3명의 사람이 있으므로 3명의 사람 위에 bbox가 그려져야 한다.

 

이제 불러온 이미지를 바탕으로  bbox를 그려본다.

 

from torchvision.ops import masks_to_boxes
# masks로 부터 boxes 생성
boxes = masks_to_boxes(masks)
print(boxes.size())
print(boxes)

>>
torch.Size([3, 4])
tensor([[ 96., 134., 181., 417.],
        [286., 113., 357., 331.],
        [363., 120., 436., 328.]])
        
from torchvision.utils import draw_bounding_boxes
drawn_boxes = draw_bounding_boxes(img, boxes, colors="red")
print(drawn_boxes)
>> 
tensor([[[ 41,  24,  29,  ...,  30,  23,  25],
         [ 38,  25,  18,  ...,  17,  13,  16],
         [ 35,  34,  16,  ...,  17,  16,  19],
         ...,
         [159, 159, 159,  ..., 183, 184, 185],
         [160, 159, 160,  ..., 184, 185, 187],
         [160, 160, 160,  ..., 184, 186, 187]],

        [[ 58,  41,  46,  ...,  53,  46,  47],
         [ 55,  42,  35,  ...,  40,  36,  38],
         [ 52,  51,  33,  ...,  37,  36,  38],
         ...,
         [159, 160, 160,  ..., 184, 187, 186],
         [158, 159, 160,  ..., 183, 186, 186],
         [158, 158, 160,  ..., 183, 185, 186]],

        [[ 42,  25,  30,  ...,  27,  20,  24],
         [ 39,  26,  19,  ...,  14,  10,  15],
         [ 36,  35,  17,  ...,  12,  11,  16],
         ...,
         [161, 162, 162,  ..., 178, 180, 180],
         [161, 161, 162,  ..., 179, 181, 182],
         [161, 161, 162,  ..., 181, 183, 182]]], dtype=torch.uint8)
      
# 위에서 선언한 show 함수를 이용하여
# 원본 이미지에 bbox가 그려진 이미지 그리기
show(drawn_boxes)

 

Figure 2.  FundanPed00054 with bbox

위와 같이 세 명의 사람 위에 bbox가 그려진 이미지의 결과를 보게된다.

원래의 tutorial에서는 아래의 과정을 거쳐서 masks를 boolean tensor로 변형한 다음,

이를 다시 masks_to_boxes로 변환하여 bbox를 그리지만 이 과정을 거치지 않아도 잘 작동함을 확인했다.

즉, mask = read_image(mask_path)로 불러온 다음 바로 masks_to_boxes(masks)를 적용해도 잘 작동한다.

 

"""
torchvision.utils.draw_bounding_boxes(
	image: Tensor, 
	boxes: Tensor, 
	labels: Optional[List[str]] = None, 
	colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None, 
	fill: Optional[bool] = False, 
	width: int = 1, 
	font: Optional[str] = None, 
	font_size: Optional[int] = None) → Tensor
"""
# labels 항목 추가
labels = ['pedestrian: 0', 'pedestrian: 1', 'pedestrian: 2']
drawn_boxes = draw_bounding_boxes(img, boxes, labels, colors="red")
show(drawn_boxes)

 

이번에는 labels를 추가하여 bbox가 어떤 라벨을 추가했는지를 확인하고자 한다.

위 코드처럼 labesl를 추가하면 결과는 아래와 같다.

 

Figure 3.  FundanPed00054 with bbox labeld

 

https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html에 나오는 example을 보다가 추가한 항목이다.

 

위 내용의 핵심은 R-CNN의 향상된 모델인 Faster R-CNN의 파인 튜닝인데 추론 파트는 아래와 같다.

 

import matplotlib.pyplot as plt

from torchvision.utils import draw_bounding_boxes, draw_segmentation_masks


image = read_image("data/PennFudanPed/PNGImages/FudanPed00046.png")
eval_transform = get_transform(train=False)

model.eval()
with torch.no_grad():
    x = eval_transform(image)
    # convert RGBA -> RGB and move to device
    x = x[:3, ...].to(device)
    predictions = model([x, ])
    pred = predictions[0]

# image normalization
image = (255.0 * (image - image.min()) / (image.max() - image.min())).to(torch.uint8)
image = image[:3, ...]

# prediction의 labels, scores, boxes를 draw_bounding_boxes에 넣은 arguments로 준비한다.
pred_labels = [f"pedestrian: {score:.3f}" for label, score in zip(pred["labels"], pred["scores"])]
pred_boxes = pred["boxes"].long() # long int로 변경한다

# 예측한 boxes, labels로 bbox를 그린다
output_image = draw_bounding_boxes(image, pred_boxes, pred_labels, colors="red")

# masks가 0.7 초과인 경우를 True로 하는 boolean masks 생성
masks = (pred["masks"] > 0.7).squeeze(1)
# Segmentation 그리기
output_image = draw_segmentation_masks(output_image, masks, alpha=0.5, colors="blue")


plt.figure(figsize=(12, 12))
plt.imshow(output_image.permute(1, 2, 0))

 

학습 과정을 생략했지만 모델은 최종적으로 output에서 boxes, labels, scores 세 가지를 리턴하게 된다.

이를 이용하여 bbox나 segmentation을 수행하면 된다.

 

 

Segmentation Example Code

Segmentation을 적용하기 위해서는 masks를 boolean tensor로 변형한다.

또한 위의 Figure 1의 그림에서도 알 수 있듯이 object는 사람 한 종류지만, 

사람이 세 명이므로 instance는 3이다.

따라서 3 x 498 x 533으로 마스크의 사이즈를 변환한다.

# We get the unique colors, as these would be the object ids.
obj_ids = torch.unique(mask)

# first id is the background, so remove it.
obj_ids = obj_ids[1:]

# split the color-encoded mask into a set of boolean masks.
# Note that this snippet would work as well if the masks were float values instead of ints.
masks = mask == obj_ids[:, None, None]

# 변형한 마스크 확인
print(masks.size())
print(masks)

>>
torch.Size([3, 498, 533])
tensor([[[False, False, False,  ..., False, False, False],
생략
 [False, False, False,  ..., False, False, False]]])

 

이제 변형한 masks를 활용하여 segmentation을 그린다.

from torchvision.utils import draw_segmentation_masks

drawn_masks = []
for mask in masks:
    drawn_masks.append(draw_segmentation_masks(img, mask, alpha=0.8, colors="blue"))

# bbox와 마찬가지로 위에서 선언한 show 함수로 segmentation을 그린다
show(drawn_masks)

 

Figure 4.  FundanPed00054 with segmentation

 

 

References:

https://pytorch.org/vision/main/auto_examples/others/plot_repurposing_annotations.html#sphx-glr-auto-examples-others-plot-repurposing-annotations-py

https://pytorch.org/vision/main/generated/torchvision.utils.draw_bounding_boxes.html

https://pytorch.org/vision/0.15/generated/torchvision.datapoints.BoundingBox.html

https://github.com/pytorch/vision/blob/main/torchvision/models/detection/roi_heads.py#L492

https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html