Torchvision object detection finetuning tutorial
이 튜토리얼에선 pre-trained된 Mask R-CNN을 finetuning한다.
데이터셋은 보행자 detection, segmentation을 위한 Penn-Fudan database를 사용한다.
345개 보행자 인스턴스가 있는 170개 이미지로 구성되어있다.
데이터셋 만들기
데이터셋은 torch.utils.data.Dataset클래스를 상속해서 __len__과 __getitem__을 구현해야한다.
__getitem__은 다음을 반환한다.
- image : height, width 사이즈의 PIL Image (width, height가 아님을 주의해야한다.)
- target : dictionary
- boxes : N개의 bounding boxes의 좌표[x0, y0, x1, y1]
- labels : 각 bounding box의 label
- image_id : 데이터셋 이미지들의 고유 id
- area : bounding box의 면적
- iscrowd : True인 경우, 평가 과정에서 제외한다.
- masks : N개의 객체마다의 분할 마스크 정보
- keypoints : N개의 객체마다의 키포인트 정보[x, y, visibility]이다. visibility=0이면 키포인트가 보이지 않는다는 것을 말한다.
여기서 사용하는 모델은 배경을 class 0으로 사용한다. 만약 준비한 데이터셋에 배경이 없다면 class 0도 없어야한다. 예를 들어, 고양이와 강아지만 분류한다면 0이 아닌 1은 고양이, 2는 강아지가 되도록 정의해야한다.
PennFudn custom dataset 작성하기
dataset download : https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip
데이터셋을 다운받은 후 압축을 푼다.
PennFudnPed dataset은 원본 이미지 PNGImages 폴더와 segmentation mask 폴더인 PedMasks 그리고 각 라벨 정보다 들어있는 Annotation 폴더로 구성되어있다.
이제 데이터셋 클래스를 작성한다.
import os
import numpy as np
import torch
from PIL import Image
class PennFudanDataset(object):
def __init__(self, root, transforms):
self.root = root
self.transforms = transforms
# 각 폴더에 해당하는 이미지 파일들을 읽어서 정렬한다.
self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages"))))
self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks"))))
def __getitem__(self, idx):
img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
img = Image.open(img_path).convert("RGB")
# 분할 마스크는 RGB로 변환하지 않는다.
mask = Image.open(mask_path)
mask = np.array(mask)
obj_ids = np.unique(mask) # 인스턴스들은 다른 색들로 인코딩 되어 있다.
obj_ids = obj_ids[1:] # 배경(class 0)은 제외한다.
masks = mask == obj_ids[:, None, None] # 컬러 인코딩된 마스크를 이진 마스크 세트로 나눈다.
num_objs = len(obj_ids)
'''
바운딩 박스 좌표 저장하는 부분
1 box = [xmin, ymin, xmax, ymax]
'''
boxes = []
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])
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.ones((num_objs,), dtype=torch.int64)
masks = torch.as_tensor(masks, dtype=torch.uint8)
image_id = torch.tensor([idx])
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
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)
return img, target
def __len__(self):
return len(self.imgs)Model 만들기
앞서 말한 대로, 이 튜토리얼에서는 Mask R-CNN을 사용한다.
Faster R-CNN은 이미지에 존재할 수 있는 객체에 대한 바운딩 박스와 클래스 점수를 모두 예측하는 모델이다. Mask R-CNN은 각 인스턴스에 대한 분할 마스크를 예측하는 추가 레이어를 Faster R-CNN에 추가한 모델이다.
torchvision modelzoo에서 사용 가능한 모델들 중 하나를 사용해서 모델을 수정하려면 보통 두가지 상황이 있다.
- 미리 학습된 모델에서 시작해서 마지막 레이어 수준만 미세 조정한다.
import torchvision from torchvision.models.detection.faster_rcnn import FastRCNNPredictor # pre-trained model load model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) # 새롭게 classifier 설정 num_classes = 2 in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
- 모델의 백본을 다른 백본으로 교체한다.
import torchvision from torchvision.models.detection import FasterRCNN from torchvision.models.detection.rpn import AnchorGenerator # pre-trained model을 읽고 features만 변수에 저장한다. backbone = torchvision.models.mobilenet_v2(pretrained=True).features backbone.out_channels = 1280 # RPN이 5(32,64,128,256,512)개의 서로 다른 크기와 3(0.5,1.0,2.0)개의 다른 측면 비율을 가진 5*3개의 anchor를 생성하도록 한다. # 각 특징 맵이 잠재적으로 다른 사이즈와 측면 비율을 가질 수 있기 때문에 Tuple[Tuple[int]] 타입을 갖도록 한다. anchor_generator = AnchorGenerator(sizes=((32,64,128,256,512),), aspect_ratios=((0.5,1.0,2.0),)) # 영역의 관심부분을 cropping 하는 feature map을 정의한다. roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], output_size=7, sampling_ratio=2) # 위 조각들을 Faster RCNN 모델로 합친다. model = FasterRCNN(backbone, num_classes=2, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler)
PennFudan dataset을 위한 instance segementation model
dataset 크기가 매우 작기 때문에, 1번 접근법을 적용한다는 점을 고려해 pre-trained model에서 미세 조정하는 방식으로 진행한다.
여기서 instance segmentation mask도 계산해야하기 때문에 Mask R-CNN을 사용한다.
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
def get_instance_segmentation(num_classes):
# pre-trained model 불러오기
model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# mask classifier를 위한 input features의 차원을 얻는다.
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
# mask predictor를 교체한다.
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)
return model전부 합치기
reference/detection/폴더에 검출 모델들의 학습과 평가를 쉽게 하기 위한 함수들이 있다.
references/detection/engine.py, references/detection/utils.py, references/detection/transforms.py를 사용한다.
위 파일들은 torchvision repo를 다운로드해서 사용한다.
%%shell
# Download TorchVision repo to use some files from
# references/detection
git clone https://github.com/pytorch/vision.git
cd vision
git checkout v0.3.0
cp references/detection/utils.py ../
cp references/detection/transforms.py ../
cp references/detection/coco_eval.py ../
cp references/detection/engine.py ../
cp references/detection/coco_utils.py ../먼저, data augmentation과 transformation 하는 코드를 작성한다.
import transforms as T
def get_transform(train):
transforms = []
transforms.append(T.ToTensor())
if train:
transforms.append(T.RandomHorizontalFlip(0.5)) # 50% 확률로 horizontal Flip한다.
return T.Compose(transforms)이제 학습과 평가를 수행하는 main function을 작성한다.
from engine import train_one_epoch, evaluate
import utils
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
num_classes = 2 # 배경(0)과 사람(1)
dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))
# 랜덤한 indices 구성해서 dataset 분할
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
datset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
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)
model = get_instance_segmentation_model(num_classes)
# GPU 사용
model.to(device)
params = [p for p in model.parameters() if p.requires_grad]
# optimizer 구성
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
# learning rate scheduler 설정
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
num_epochs = 10
for epoch in range(num_epochs):
train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
lr_scheduler.step()
evaluate(model, data_loader_test, device=device)
print("That's it!")학습 중 1 epoch이 진행되면 아래와 같이 COCO 스타일의 mAP를 얻을 수 있다.
IoU threshold range를 0.5에서 0.95까지 0.05 step으로 증가시켜 AP를 계산함.
Epoch: [0] [ 0/60] eta: 0:01:08 lr: 0.000090 loss: 4.4058 (4.4058) loss_classifier: 0.7733 (0.7733) loss_box_reg: 0.4066 (0.4066) loss_mask: 3.2090 (3.2090) loss_objectness: 0.0109 (0.0109) loss_rpn_box_reg: 0.0060 (0.0060) time: 1.1382 data: 0.3330 max mem: 3495
Epoch: [0] [10/60] eta: 0:00:33 lr: 0.000936 loss: 1.5725 (2.2352) loss_classifier: 0.4764 (0.4605) loss_box_reg: 0.2614 (0.2749) loss_mask: 0.8914 (1.4729) loss_objectness: 0.0178 (0.0217) loss_rpn_box_reg: 0.0056 (0.0052) time: 0.6762 data: 0.0368 max mem: 3733
Epoch: [0] [20/60] eta: 0:00:26 lr: 0.001783 loss: 1.1325 (1.6025) loss_classifier: 0.2917 (0.3674) loss_box_reg: 0.2613 (0.2794) loss_mask: 0.4333 (0.9290) loss_objectness: 0.0178 (0.0185) loss_rpn_box_reg: 0.0067 (0.0082) time: 0.6357 data: 0.0077 max mem: 3733
Epoch: [0] [30/60] eta: 0:00:19 lr: 0.002629 loss: 0.6684 (1.2818) loss_classifier: 0.1261 (0.2837) loss_box_reg: 0.2289 (0.2661) loss_mask: 0.2537 (0.7074) loss_objectness: 0.0111 (0.0165) loss_rpn_box_reg: 0.0090 (0.0081) time: 0.6246 data: 0.0077 max mem: 3733
Epoch: [0] [40/60] eta: 0:00:12 lr: 0.003476 loss: 0.4744 (1.0749) loss_classifier: 0.0770 (0.2303) loss_box_reg: 0.1879 (0.2480) loss_mask: 0.1981 (0.5757) loss_objectness: 0.0051 (0.0135) loss_rpn_box_reg: 0.0061 (0.0075) time: 0.6053 data: 0.0077 max mem: 3733
Epoch: [0] [50/60] eta: 0:00:06 lr: 0.004323 loss: 0.3667 (0.9306) loss_classifier: 0.0533 (0.1939) loss_box_reg: 0.1454 (0.2235) loss_mask: 0.1537 (0.4945) loss_objectness: 0.0016 (0.0113) loss_rpn_box_reg: 0.0041 (0.0073) time: 0.6172 data: 0.0078 max mem: 3733
Epoch: [0] [59/60] eta: 0:00:00 lr: 0.005000 loss: 0.3054 (0.8390) loss_classifier: 0.0383 (0.1709) loss_box_reg: 0.0994 (0.2058) loss_mask: 0.1448 (0.4452) loss_objectness: 0.0020 (0.0101) loss_rpn_box_reg: 0.0041 (0.0070) time: 0.6153 data: 0.0073 max mem: 3747
Epoch: [0] Total time: 0:00:37 (0.6289 s / it)
creating index...
index created!
Test: [ 0/170] eta: 0:01:06 model_time: 0.1671 (0.1671) evaluator_time: 0.0049 (0.0049) time: 0.3887 data: 0.2141 max mem: 3747
Test: [100/170] eta: 0:00:09 model_time: 0.1223 (0.1156) evaluator_time: 0.0082 (0.0079) time: 0.1447 data: 0.0040 max mem: 3747
Test: [169/170] eta: 0:00:00 model_time: 0.1060 (0.1158) evaluator_time: 0.0039 (0.0074) time: 0.1193 data: 0.0041 max mem: 3747
Test: Total time: 0:00:22 (0.1302 s / it)
Averaged stats: model_time: 0.1060 (0.1158) evaluator_time: 0.0039 (0.0074)
Accumulating evaluation results...
DONE (t=0.03s).
Accumulating evaluation results...
DONE (t=0.03s).
IoU metric: bbox
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.658
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.982
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.853
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.279
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.573
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.663
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.297
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.726
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.726
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.520
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.723
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.729
IoU metric: segm
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.731
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.982
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.905
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.187
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.449
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.741
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.322
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.770
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.771
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.500
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.732
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.778아래 첫번째 이미지에 대한 instance segmentation mask 결과를 확인할 수 있다.
마무리
이 튜토리얼에서는 custom dataset을 가지고 instance segmentation model을 위한 학습 파이프라인을 만드는 방법에 대해 배웠다.
데이터셋을 생성하기 위해 작성했던 torch.utils.data.Dataset은 images와 ground truth boxes, segmentation masks를 반환한다.
COCO train 2017로 pre-trained된 Mask R-CNN을 transfer learning 하였다.
Uploaded by Notion2Tistory v1.1.0