[Pytorch] PASCAL VOC 2012 Segmentation Dataset 사용하기

안녕하세요. 오늘은 영상 분할 (Image Segmentation) 분야에서 필수적으로 사용되는 데이터셋인 PASCAL VOC 2012 Segmentation에 대해서 알아보고 pytorch를 이용해서 DataLoader를 만들어보도록 하겠습니다.

 

1. PASCAL VOC 2012 Segmentation Dataset 소개

PASCAL VOC 2012 데이터셋은 Semantic Segmentation 뿐만 아니라 Object Detection, Instance Segmentation을 위한 다용도 데이터셋입니다. 저는 이 중에서 Semantic Segmentation을 위한 DataLoader를 구성해보도록 하겠습니다. 해당 데이터셋은 21개의 클래스 (vehicles, household, animals, and other: aeroplane, bicycle, boat, bus, car, motorbike, train, bottle, chair, dining table, potted plant, sofa, TV/monitor, bird, cat, cow, dog, horse, sheep, and person)가 pixel-level 에서 label이 부여된 데이터셋입니다. 해당 데이터셋은 1,464 장의 훈련 데이터셋과 1,449 장의 검증 데이터셋이 포함되어 있습니다. 하지만, 시험 데이터셋은 지난 포스팅의 ImageNet과 마찬가지로 비공개 상태 입니다. 따라서, 실질적으로 실험 및 연구를 진행할 때는 검증 데이터를 시험 데이터셋으로 사용하고 훈련 데이터셋을 훈련 데이터셋과 검증 데이터셋으로 나누면 됩니다. 

 

2. PASCAL VOC 2012 Segmentation DataLoader

오늘 구현할 DataLoader는 주로 딥러닝 공부방의 게시글을 참조하였음을 먼저 알려드립니다. 또한, DataLoader를 구현한 전체 코드는 제 깃허브에 있습니다. 

 

GitHub - skawngus1111/AwesomeDeepLearning

 

COCO 데이터셋이나 ImageNet의 경우에는 저희가 직접 데이터셋을 다운받아야하지만 PASCAL VOC의 경우에는 pytorch에서 제공하는 VOCSegmentation 데이터셋을 이용하면 자동으로 다운로드까지 받아주기 때문에 쉽게 구현할 수 있습니다. 

 

class VOCSegDataset(VOCSegmentation):
    def __getitem__(self, idx):
        image = Image.open(self.images[idx]).convert('RGB')
        label = Image.open(self.masks[idx])

        if self.transforms is not None:
            seed = random.randint(0, 2 ** 32)
            self._set_seed(seed); image = self.transforms(image)
            self._set_seed(seed); label = self.transforms(label) * 255
            label[label > 20] = 0

        return image, label

    def _set_seed(self, seed):
        random.seed(seed)
        torch.manual_seed(seed)

 

원래 Dataset 클래스를 정의할 때는 생성자 함수 (__init__)를 정의하지만 VOCSegmentation을 상속받음으로써 훨씬 간단하게 __getitem__ 함수만 구현하면 됩니다. 

 

train_ds = VOCSegDataset(DATA_DIR, year='2012', image_set='train', download=True, transforms=image_transforms)
val_ds = VOCSegDataset(DATA_DIR, year='2012', image_set='val', download=True, transforms=image_transforms)

 

이제 위와 같이 year, image_set을 선택하면 2012년도의 train/val 데이터셋을 자동으로 다운로드 해줍니다. 참으로 쉽죠? 

 

for img, target in train_loader:
    plt.imshow(np.transpose(img[0].cpu().detach().numpy(), (1, 2, 0)))
    plt.show()
    plt.imshow(target[0].squeeze().cpu().detach().numpy())
    plt.show()

 

그리고 위와 같이 반복문을 구성하여 데이터를 로딩하면 쉽게 시각화까지 진행할 수 있습니다. 다만, 주의하실 점은 제 코드를 그대로 사용하시면 image transformation 단계에서 segmentation label이 0~20 이여야하는데 255로 강제로 나누어지는 문제가 있습니다. 이를 해결하고자 255를 곱하는 과정을 추가하였습니다. 혹시, 이 방법말고 다른 좋은 방법을 아시는 분은 댓글로 남겨주시면 감사하겠습니다. ㅎㅎ

 

마무리로 PASCAL VOC 2012의 Segmentation 데이터셋의 시각화를 진행해보도록 하겠습니다. 

 

 

참고문헌

[1]. https://paperswithcode.com/dataset/pascal-voc

Papers with Code - PASCAL VOC Dataset

[2]. https://deep-learning-study.tistory.com/707

[PyTorch] VOC Segmentation 데이터셋 사용하기