09 객체 탐지 (6) Mask R-CNN pytorch 실습

이번 포스팅에서 사용하는 데이터는 책에서 제공하는 데이터로 간단히 설명하면 MS COCO 데이터를 개와 고양이에 대해서만 소규모로 샘플링한 데이터이다.

 

이전 Faster R-CNN 실습 코드를 일부 수정하여 Mask R-CNN 모델을 학습해 본다. 여기서는 수정되는 부분만 설명하기 때문에 이해가되지 않는 코드가 있다면 이전 포스팅을 참고하길 바란다.

2024.02.20 - [책/파이토치 트랜스포머를 활용한 자연어 처리와 컴퓨터비전 심층학습] - 09 객체 탐지 (2) Faster R-CNN pytorch 실습

09 객체 탐지 (2) Faster R-CNN pytorch 실습

 

 

import os
import numpy as np
from pathlib import Path

import matplotlib.pyplot as plt
from PIL import Image

import torch
from torch import optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.transforms.functional import to_pil_image
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

from pycocotools import mask as maskUtils
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

 

데이터셋

이전과 마찬가지로 pycocotools.coco의 COCO API로 json 형식의 annotation 파일을 불러온다.

root = DATA_PATH / "datasets" / "coco"
ann_path = os.path.join(root, "annotations", "train_annotations.json")
image_path = os.path.join(root, "train")
coco = COCO(ann_path)

 

MS COCO 데이터셋의 segmentation 정보는 다음과 같이 폴리곤 형태이다.

import matplotlib.patches as patches

img_id = 272153
file_name = coco.loadImgs(img_id)[0]["file_name"]
image = Image.open(os.path.join(image_path, file_name))

anns = coco.loadAnns(coco.getAnnIds(img_id))
segmentations = anns[0]["segmentation"]

# segmentation 정보를 numpy 배열로 변환
segmentation_np = np.array(segmentations)

# x, y 좌표를 분리
x = segmentation_np[0, ::2]
y = segmentation_np[0, 1::2]

# 그림 생성
fig, ax = plt.subplots()
ax.imshow(image)

# 다각형 그리기
shp = patches.Polygon(np.column_stack((x, y)), edgecolor='r', facecolor='none')
plt.gca().add_patch(shp)
plt.axis('scaled')

plt.show()

 

이를 모델 학습에 사용하기 위해서는 이진 마스크 형태로 변경해주어야 한다. 이를 위해 pycocotools의 frPyObjects API를 활용한다.

 

이는 segmentation 정보를 RLE(Run-Length Encoding) 형식으로 변환하는 역할을 한다. RLE 형식의 정보는 후에 maskUtils.decode를 통해 이진 마스크로 디코딩된다.

 

width, height = image.size

binary_mask = []
for seg in segmentations:
  rles = maskUtils.frPyObjects([seg], height, width)
  binary_mask.append(maskUtils.decode(rles))

combined_mask = np.sum(binary_mask, axis=0).squeeze()
fig, ax = plt.subplots(1, 2, figsize=(14, 14))
ax[0].imshow(image)
ax[1].imshow(combined_mask)
plt.axis("off")
plt.show()

 

 

이를 활용해 데이터셋 클래스를 정의하고 데이터로더를 생성한다.

class COCODataset(Dataset):
  def __init__(self, root, train, transform=None):
    super().__init__()
    directory = "train" if train else "val"
    annotations = os.path.join(root, "annotations", f"{directory}_annotations.json")

    self.coco = COCO(annotations)
    self.image_path = os.path.join(root, directory)
    self.transform = transform

    self.categories = self._get_categories()
    self.data = self._load_data()
  
  def _get_categories(self):
    categories = {0: "background"}
    for category in self.coco.cats.values():
      categories[category["id"]] = category["name"]
    return categories
  
  def _load_data(self):
    data = []
    for _id in self.coco.imgs:
      file_name = self.coco.loadImgs(_id)[0]["file_name"]
      image_path = os.path.join(self.image_path, file_name)
      image = Image.open(image_path).convert("RGB")
      width, height = image.size

      boxes, labels, masks = [], [], []
      anns = self.coco.loadAnns(self.coco.getAnnIds(_id))
      for ann in anns:
        x, y, w, h = ann["bbox"]
        segmentations = ann["segmentation"]
        try:
          mask = self._polygon_to_mask(segmentations, width, height)
        except:
          pass
        
        boxes.append([x, y, x + w, y + h])
        labels.append(ann["category_id"])
        masks.append(mask)
      
      target = {
          "image_id": torch.LongTensor([_id]),
          "boxes": torch.FloatTensor(boxes),
          "labels": torch.LongTensor(labels),
          "masks": torch.FloatTensor(masks),
      }
      data.append([image, target])
    return data
  
  def _polygon_to_mask(self, segmentations, width, height):
    binary_mask = []
    for seg in segmentations:
      rles = maskUtils.frPyObjects([seg], height, width)
      binary_mask.append(maskUtils.decode(rles))
    
    combined_mask = np.sum(binary_mask, axis=0).squeeze()
    return combined_mask
  
  def __getitem__(self, index):
    image, target = self.data[index]
    if self.transform:
      image = self.transform(image)
    return image, target
  
  def __len__(self):
    return len(self.data)

def collator(batch):
    return tuple(zip(*batch))

transform = transforms.Compose(
    [
        transforms.PILToTensor(),
        transforms.ConvertImageDtype(dtype=torch.float)
    ]
)

train_dataset = COCODataset(root, train=True, transform=transform)
test_dataset = COCODataset(root, train=False, transform=transform)

train_dataloader = DataLoader(
    train_dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=collator
)
test_dataloader = DataLoader(
    test_dataset, batch_size=1, shuffle=True, drop_last=True, collate_fn=collator
)

 

 

모델 정의

• 클래스 개수는 배경, 개, 고양이로 3개이다.
• hidden_layer는 마스크 예측을 위한 중간 특징 맵의 차원을 의미한다. 여기서는 ResNet-50을 사용하고 있는데 ResNet-50의 스테이지 1의 출력 차원이 256이기 때문에 다음과 같이 설정한다.
• 모델은 torchvision에서 제공하는 사전 학습된 Mask R-CNN 모델을 사용한다. 우리는 3개의 클래스를 예측하기 위한 모델을 정의하기 때문에 model의 box_predictor와 mask_predictor를 새롭게 지정해준다.

num_classes = 3
hidden_layer = 256

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = maskrcnn_resnet50_fpn(weights="DEFAULT")

model.roi_heads.box_predictor = FastRCNNPredictor(
    in_channels=model.roi_heads.box_predictor.cls_score.in_features,
    num_classes=num_classes,
)

model.roi_heads.mask_predictor = MaskRCNNPredictor(
    in_channels=model.roi_heads.mask_predictor.conv5_mask.in_channels,
    dim_reduced=hidden_layer,
    num_classes=num_classes,
)
model.to(device)

 

 

학습

loss_dict는 딕셔너리 구조이며 손실값으로 loss_classifier, loss_box_reg, loss_mask, loss_objectness, loss_rpn_box_reg이 있다.

params = [p for p in model.parameters() if p.requires_grad]
optimizer = optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)

for epoch in range(5):
    cost = 0.0
    for idx, (images, targets) in enumerate(tqdm(train_dataloader, desc=f'Epoch {epoch+1}/{5}', unit='batch')):
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

        cost += losses

    lr_scheduler.step()
    cost = cost / len(train_dataloader)
    print(f"Epoch : {epoch+1:4d}, Cost : {cost:.3f}")

Epoch 1/5: 100%|██████████| 607/607 [04:18<00:00,  2.35batch/s]
Epoch :    1, Cost : 0.418
Epoch 2/5: 100%|██████████| 607/607 [04:16<00:00,  2.37batch/s]
Epoch :    2, Cost : 0.291
Epoch 3/5: 100%|██████████| 607/607 [04:18<00:00,  2.35batch/s]
Epoch :    3, Cost : 0.250
Epoch 4/5: 100%|██████████| 607/607 [04:19<00:00,  2.34batch/s]
Epoch :    4, Cost : 0.222
Epoch 5/5: 100%|██████████| 607/607 [04:23<00:00,  2.31batch/s]Epoch :    5, Cost : 0.205

 

 

추론 결과

 

 

추론 결과를 보면 Faster R-CNN에 비해 작은 객체도 정말 잘 찾아낸다. 개인적으로 조금 놀라웠다. 두 번째 그림처럼 코끼리를 개로 예측하기도 하는데 이는 우리가 개와 고양이만 있는 데이터만 사용하면 학습하였기 때문에 어쩔 수 없는 부분이고, 전체 데이터셋으로 학습한다면 성능이 더 좋을 것이다.

 

 

평가

segmentation mask 평가는 예측과 정답 마스크를 0과 1로 구성된 이진 마스크로 변환한 뒤 둘의 교집합과 합집합을 계산한다. 그런 다음 (교집합 / 합집합)으로 마스크 IoU를 계산한다.

 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.558
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.803
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.664
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.188
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.607
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.566
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.603
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.694
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.694
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.235
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.684
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.719

 

전체적인 성능은 첫 번째 줄에서 확인할 수 있는데 Faster R-CNN이 0.262였던 것에 비해 0.558로 성능이 많이 좋아졌다. 또한 네 번째 줄은 작은 객체에 대한 성능 평가를 나타내는데 이 또한 0.005에서 0.188로 많이 상승했다.