[KUBIG; Basic-CV] TorchVision Object Detection

 

 

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

TorchVision Object Detection Finetuning Tutorial

Pretrained된 Mask R-CNN 모델을 Penn-Fudan Database for Pedestrian Detection and Segmentation 데이터셋에 맞게 미세 조정(finetuning)하는 방법을 설명

1. Dataset Class 정의

torch.utils.data.Dataset를 상속받는 사용자 정의 데이터셋 클래스는 image, bounding box, label, mask를 포함하는 target dictionary를 반환해야 함.

▶ Code

import os
import numpy as np
import torch
from PIL import Image

# PennFudanDataset 클래스 정의, torch.utils.data.Dataset를 상속
class PennFudanDataset(torch.utils.data.Dataset):
    def __init__(self, root, transforms):
        self.root = root  # 데이터셋 루트 디렉토리
        self.transforms = transforms  # img transform
        self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages"))))  # image 파일 list
        self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks"))))  # mask 파일 list

    def __getitem__(self, idx):
        img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])  # image 경로
        mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])  # mask 경로
        img = Image.open(img_path).convert("RGB")  # image 열기 및 RGB로 변환
        mask = Image.open(mask_path)  # mask image 열기
        mask = np.array(mask)  # mask를 numpy 배열로 변환
        obj_ids = np.unique(mask)  # 객체 ID 추출
        obj_ids = obj_ids[1:]  # 배경 ID 제외

        masks = mask == obj_ids[:, None, None]  # 각 객체에 대한 mask 생성
        num_objs = len(obj_ids)  # 객체 수
        boxes = []  # bounding box 리스트 초기화
        for i in range(num_objs):
            pos = np.where(masks[i])  # 객체 위치 찾기
            xmin = np.min(pos[1])
            xmax = np.max(pos[1])
            ymin = np.min(pos[0])
            ymax = np.max(pos[0])
            boxes.append([xmin, ymin, xmax, ymax])  # bounding box 추가

        boxes = torch.as_tensor(boxes, dtype=torch.float32)  # bounding box를 텐서로 변환
        labels = torch.ones((num_objs,), dtype=torch.int64)  # 객체 label (모두 1로 설정)
        masks = torch.as_tensor(masks, dtype=torch.uint8)  # mask를 텐서로 변환

        image_id = torch.tensor([idx])  # 이미지 ID 텐서
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])  # 각 bounding box의 면적
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)  # 객체가 군집인지 여부 (모두 0으로 설정)

        # target 딕셔너리 생성
        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        if self.transforms is not None:
            img, target = self.transforms(img, target)  # transform 적용

        return img, target  # image와 target 반환

    def __len__(self):
        return len(self.imgs)  # 데이터셋 크기 반환

 

2. 모델 정의

Mask R-CNN을 사용하며, 데이터셋의 클래스 수에 맞게 분류기를 교체

▶ Mask R-CNN

1. Background

 

• classification: 이미지의 클래스를 판별
• localization: 객체의 위치정보를 추정
• object detection: 위치와 클래스 정보 모두 반환
• semantic segmentation: bounding box 대신 객체 주변에 경계를 그리고 픽셀 수준의 객체 정보 반환
• instance segmentation: semantic segmentation와 다른 점은 동일한 클래스에 대해서도 서로 다른 instance로 인식한다는 것

2. Faster R-CNN Framework

• Deep ConvNet: 입력 이미지를 CNN에 입력하여 feature map 추출
• RoI Projection: Region of Interest (RoI) Projection을 통해 특징 맵에서 관심 영역을 추출
• RoI pooling layer: RoI pooling 레이어는 관심 영역을 고정된 크기의 벡터로 변환
• FCs (Fully Connected Layers): Fully Connected Layers을 통해 RoI 특징 벡터를 처리
• Output (bbox regressor, softmax): Bounding box regressor은 객체의 bounding box 예측, softmax는 각 RoI가 특정 클래스에 속할 확률 예측

3. Mask R-CNN Framework

 

Faster R-CNN은 Feature Extraction과 ROI proposal을 위한 모든 기초 작업을 구축하고 Mask R-CNN은 이 모델에 편승하여 ROI pooling 후에 convolution layer 2개를 추가하여 Mask를 추출하는 방법.

Network Architecture

Backbone 이미지의 feature을 추출하기 위해 ResNet과 ResNeXt의 50, 101 layer과 FPN(Feature Pyramid Network)을 backbone으로 사용하고, Faster R-CNN의 Head(Classification and Regression)에 Mask branch를 추가하며 backbone (ResNet, FPN)에 따라 Head의 구조가 달라짐

 

이때 추가되는 두 layer의 backbone은 FPN(Feature Pyramid Network) 심층 신경망이며, 이는 상향식 경로(bottom-up path), 하향식 경로(top-down path), 측면 연결(lateral connection)로 구성.

• 상향식 경로 (Bottom-up Path): 원래 이미지로부터 특징을 추출하는 단계 (고 → 저해상도)로, ResNet이나 VGG와 같은 ConvNet이 사용되며 이 단계에서는 이미지의 다양한 수준의 특징들이 추출
• 하향식 경로 (Top-down Path): 상향식 경로에서 추출된 특징 맵을 입력으로 받아, 이와 비슷한 크기의 특징 피라미드 맵을 생성 (저 → 고해상도) => 상향식 경로에서 놓칠 수 있는 세부 정보를 보완
• 측면 연결 (Lateral Connection): 상향식 경로와 하향식 경로 사이의 각 레벨을 연결하여 정보를 통합. 컨볼루션과 추가 연산을 통해 이루어지며, 두 경로의 장점을 결합하여 더 강력한 특징 표현을 제공

cf) 고해상도 feature map: 이미지의 세부 정보 많이 포함, 저해상도 feature map: 이미지의 전반적인 구조와 큰 물체를 잘 포착

위 과정을 거친 후, FPN에서 추출된 특징 맵을 처리하는 첫 번째 단계로 RPN(Region Proposal Network)을 통해 객체가 포함될 가능성이 있는 영역(Region Proposals)을 앵커(anchor)를 사용하여 원시 이미지의 다양한 위치와 스케일에서 객체 후보를 제안. 이때, 앵커는 이미지에 대해 미리 정의된 위치와 배율을 가진 상자 집합으로 객체 또는 배경으로 분류. IoU (Intersection over Union) 값을 기준으로 앵커와 실제 객체 경계 상자의 겹침 정도를 계산하여 앵커를 할당. (High IoU는 객체를 포함하는 앵커로 할당(양성 앵커), Low IoU는 배경으로 간주(음성 앵커))

 

두번째 단계는 ROI Align으로, 첫 번째 단계(RPN)에서 제안된 영역을 받아서 객체 클래스, 경계 상자, 마스크를 생성. 앵커 없이 특징 맵에서 관련 영역을 찾고, 각 객체에 대한 마스크를 픽셀 단위로 생성하는 방법. 기존의 ROI pooling은 디지털화된 warping으로 대상 셀의 경계를 원본 특징 맵의 셀 경계에 맞추기 위해 강제로 조정하여 정보 손실이 발생할 수 있으며, 셀 크기가 균일하지 않아 feature map에서의 정확한 위치 정보가 손실될 수 있지만, ROI Align은 셀 경계를 디지털화하지 않고, 대상 셀의 크기를 일정하게 유지하며 각 셀 내의 feature map 값을 더 정확히 계산하기 위해 interpolation을 적용하여 더 정밀한 값을 산출.

(좌) ROI pooling, (우) ROI align

 

Summary) input image 내 객체에 대한 bounding box 생성 → Region of Interest Align (RoIAlign)은 특징 맵에서 고정된 크기의 피처 벡터를 추출 → 추출된 피처 벡터를 여러 개의 convolution layer를 통과시켜 처리 → 최종적으로 각 객체의 클래스와 바운딩 박스를 예측하고, 객체의 세그멘테이션 마스크를 생성

 

 

 

▶ Code

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

def get_model_instance_segmentation(num_classes):
    # pre-trained된 Mask R-CNN 모델 불러오기
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)

    # Classifier head의 input feature 수
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # pre-trained head 교체 (num_classes는 데이터셋의 클래스 수)
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # Mask classifer의 input feature 수
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256  # 중간 layer의 크기
    # 새로운 mask predictor로 교체
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)

    return model  # fine tuning된 모델을 반환

 

3. 훈련 및 평가

사용자 정의 데이터셋을 사용하여 모델을 훈련하고 성능을 평가

▶ Code

import torch
from engine import train_one_epoch, evaluate
import utils
import transforms as T

# img tranform 함수 정의
def get_transform(train):
    transforms = []
    transforms.append(T.ToTensor())  # image를 텐서로 변환
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))  # 학습 시 랜덤으로 좌우 반전
    return T.Compose(transforms)

# PennFudanDataset class instance 생성 (훈련 및 테스트용)
dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))

# data loader 설정 (훈련 및 테스트용)
data_loader = torch.utils.data.DataLoader(
    dataset, batch_size=2, shuffle=True, num_workers=4,
    collate_fn=utils.collate_fn)

data_loader_test = torch.utils.data.DataLoader(
    dataset_test, batch_size=1, shuffle=False, num_workers=4,
    collate_fn=utils.collate_fn)

# 디바이스 설정
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# 클래스 수 설정 (배경 + 객체)
num_classes = 2
# 모델 인스턴스 생성 및 디바이스로 이동
model = get_model_instance_segmentation(num_classes)
model.to(device)

# 옵티마이저 및 learning rate scheduler 설정
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

num_epochs = 10  # 학습 epoch 수 설정

# 학습 및 평가 loop
for epoch in range(num_epochs):
    # 한 epoch 동안 모델 학습
    train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
    # learning rate scheduler 업데이트
    lr_scheduler.step()
    # 모델 평가
    evaluate(model, data_loader_test, device=device)