[Penn-Fudan] Object detection 학습하기 - 1

 - 본 포스팅은 Pytorch 홈페이지에서 제공하는 예제를 기반으로 작성하였습니다.

 

이번 포스팅에서는 Mask R-CNN을 사용하여 Object detection을 수행해보자.

정확히는 처음부터 학습하는 것은 아니고, 이미 학습된 모델을 fine-tunning하는 형식으로 진행할 것이다.

 

주요 절차는 아래와 같다.

 

1. 데이터셋 다운로드 및 정의
2. 학습 모델 정의
3. 모델 성능 검증

---

1. 데이터셋 다운로드 및 정의

 

먼저 학습에 필요한 데이터셋을 다운받아야한다. 아래의 링크에 접속하여 4. Download zipoed file를 눌러서 다운받자.

 

https://www.cis.upenn.edu/~jshi/ped_html/

Pedestrian Detection Database

 

받은 파일을 압축 해제하면, 아래의 사진처럼 이미지들과 annotation을 확인할 수 있다.

 

 

잠시 Penn-Fudan 데이터셋의 구조에 대해 설명하면, 이름과 사진을 통해 알 수 있듯이 보행자에 대한 이미지 데이터셋이다.

 

Penn-Fudan Datset 예시 사진

 

위의 예시처럼 총 170개의 이미지와 마스크 이미지 쌍으로 구성되어있으며, 다양한 이미지 스케일을 갖고있다.
그리고 전체 데이터셋에는 345개의 보행자 인스턴스가 존재한다.

 

이제 저장된 이미지 파일들을 사용하기 위한 데이터셋을 정의해보자. 코드는 다음과 같다.

 

import os
import torch

from torch.utils.data import Dataset
from torchvision.ops.boxes import masks_to_boxes
from torchvision import tv_tensors
from torchvision.transforms.v2 import functional as F
from torchvision.io import read_image

root = '../../Data/PennFudanPed'

class PennFudanDataset(Dataset):
    def __init__(self, transforms=None):
        self.transforms = transforms
        self.img_path = root+'/PNGImages'
        self.msk_path = root+'/PedMasks'
        self.img_list = sorted(os.listdir(self.img_path))
        self.msk_list = sorted(os.listdir(self.msk_path))
    
    def __getitem__(self, idx):
        img = read_image(os.path.join(self.img_path, self.img_list[idx]))
        msk = read_image(os.path.join(self.msk_path, self.msk_list[idx]))

        # Extract unique object IDs from mask image
        obj_ids = torch.unique(msk)

        # The 1st ID is background, so it is removed
        obj_ids = obj_ids[1:]
        num_objs = len(obj_ids)

        # Split the mask to binary masks according to id value
        msks = (msk == obj_ids[:, None, None]).to(dtype=torch.uint8)

        # Get bounding box coordinate for each mask
        boxes = masks_to_boxes(msks)

        # Only one class (Pedestrian) is considered, with background already removed. 
        # All remaining masks belong to the Pedestrian class.
        lbls = torch.ones((num_objs,), dtype=torch.int64)

        img_id = idx
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])

        # Suppose all instances are not crowd
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        # Wrap the img and target with torchvision tensor
        img = tv_tensors.Image(img)

        target = {}
        target["boxes"] = tv_tensors.BoundingBoxes(boxes, format="XYXY",
                                                   canvas_size = F.get_size(img))
        target["masks"] = tv_tensors.Mask(msks)
        target["labels"] = lbls
        target["image_id"] = img_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.img_list)

 

728x90

대부분의 코드는 다른 task에서 사용하는 방식과 비슷하다. 여기서 잠시 target 부분을 정리하고 넘어가자.

 

우선 target은 모델이 주어진 이미지를 통해 학습해야할 목표에 대한 내용을 담고있다. 

코드를 순서대로 살펴보면, 한 장의 마스크 이미지에 존재하는 여러 인스턴스들을 각각 분석한다.

 

그 후, 배경값을 제외하고  인스턴스 값에 따라 각각 새로운 마스크 이미지들로 만들어주고 이를 msks로 저장해준다.

따라서 msks는 (instance 수, image Height, image Width)와 같은 구성을 갖게될 것이다.

 

다음은 masks_to_boxes 함수를 통해 각 인스턴스의 bounding box 위치를 찾아준다.

각 인스턴스마다 (x_min, y_min, x_max, y_max)의 값을 저장하게된다.

 

다음은 각 인스턴스가 어떤 클래스에 속하는지 정의해준다. 이 데이터셋에서 우리는 보행자(Pedstrain : 1)인지

배경 (Background : 0)인지만 다루게 될 것이므로, obj_ids에 저장된 인스턴스들은 모두 1이 되는 것이다.

 

따라서, lbls를 torch.ones를 사용해서 인스턴스 수에 해당하는 길이로 생성해준다.

 

설명이 필요한 부분은 마무리되었고, target에 입력될 값들을 정리하면 다음과 같다.

• boxes : (N, 4) 형태. N은 인스턴스 수(=Bounding box 수)고 각 인스턴스 마다의 (x_min, y_min, x_max, y_max)개의
               정보를 담고 있다. [tv_tensors.BoundingBoxes]
• labels : 각 인스턴스가 속한 class 정보 [torch.Tensors]
• image_id : 인스턴스 id 값 [int]
• area : Bounding Box의 면적. COCO metric에서 사용 [torch.Tensor]
• iscorwd : 이미지가 군중사진(여러 사람이 뭉쳐져있는 사진)인지의 여부. [torch.Tensor]
• (optionally) masks : 인스턴스 별 마스크 이미지 (N, H, W)형태

---

 

이번 포스팅은 Penn-Fudan 데이터셋을 다운받고 정의하는 과정에 대해 알아보았다.

다음 포스팅은 학습에 사용할 모델을 정의하는 부분에 대해 다뤄보겠다.

 

질문이나 오류에 대한 답글 남겨주시면 감사하겠습니다 :)