[Dataset] PyTorch의 Dataset 활용하기

PyTorch와 Dataset

이번 포스팅에서는 학습을 위해 간단한 Dataset을 만들려고 합니다.

PyTorch의 torch.utils.data.DataLoader와 torch.utils.data.Dataset는 학습을 위한 데이터를 더욱 편리하게 관리할 수 있게 합니다.

COCO 데이터 셋을 이용하여 이전 포스팅과 내용이 이어지니 참고하시면 도움이 될 것 같습니다.

https://fictitious.tistory.com/2

[Python] Pycocotools를 이용한 COCO Dataset 확인하기

 

 

 

---

1. Class 정의

이전 포스팅에서 COCO 데이터 셋의 영상과 Annotation 정보에 접근하는 방법을 공부했습니다.

segmentation, keypoints 등 많은 정보가 있었지만 그 중 영상과 bbox만을 학습에 이용하기 위해 

MyCocoDataset이라는 객체를 아래와 같이 정의하였습니다.

 

필요한 모듈들을 import하고 torch.utils.data의 Dataset을 상속합니다.

초기화 단계에서는 우선 root와 json파일인 annoFile을 받도록 하였으며 이를 이용하여 COCO 객체(coco)를 생성합니다.

또한, 보고자하는 카테고리 인덱스(catIds)와 해당 카테고리를 포함하고 있는 영상 인덱스(imgIds)를 함께 구합니다. 

import os

import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import transforms
from PIL import Image

import matplotlib.pyplot as plt
import numpy as np

class MyCocoDataset(Dataset):
    def __init__(self, root, annoFile):
        super().__init__()
        from pycocotools.coco import COCO
        self.root = root
        self.coco = COCO(annoFile)
        self.catIds = self.coco.getCatIds(catNms=['person'])
        self.imgIds = self.coco.getImgIds(catIds = self.catIds)

 

 

 

다음으로 __getitem__(인덱스에 접근할 때 호출)부분을 작성합니다.

저는 영상과 bbox를 학습에 사용할 예정이므로 해당 정보들을 획득합니다.

마찬가지로 getAnnIds와 loadAnns를 이용하여 해당 영상의 Annotation 정보에 접근할 수 있습니다.

그 중 bbox만 추출하기 위해 bboxes라는 리스트를 만들어 Annotation의 bbox 부분(ann['bbox'])만 추가합니다.

참고로 bbox는 4개의 숫자로 구성되어 있으므로 bboxes의 크기는 (사람 수, 4)가 되게 됩니다.

 

영상을 얻기위해 loadImgs로 영상 이름을 얻고 해당 파일에 접근하여 영상을 엽니다.

학습에 사용하기 위해서 영상(img)를 tensor로 변환하며 함께 Normalize도 진행합니다.

여기서 영상으로 출력할 수 없으므로 꼭 tensor로 변환하셔야 합니다.

class CocoDataset(Dataset):
    
    def __getitem__(self, index): # 인덱스에 접근할 때 호출
        imageId = self.imgIds[index]
        annIds = self.coco.getAnnIds(imgIds = imageId, catIds=self.catIds)
        anns = self.coco.loadAnns(annIds)
        
        bboxes = []
        for ann in anns:
            bbox = np.array(ann['bbox'])
            bboxes.append(bbox)
        bboxes = np.array(bboxes)
        
        imageInfo = self.coco.loadImgs(imageId)[0]
        with open(os.path.join(self.root, 'images', 'val2014', imageInfo['file_name']), 'rb') as f:
            img = Image.open(f).convert('RGB')
            
        preprocess = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                     std=[0.1, 0.1, 0.1]),
        ])
        img = preprocess(img)
        return img, bboxes
    
    def __len__(self):
        return len(self.imgIds)

 

 

 

---

2. DataLoader

이번에는 앞서 정의한 데이터 셋 객체를 생성하여 torch.utils.data.DataLoader로 불러오겠습니다.

필요한 변수들을 작성하여 MyCocoDataset 객체를 생성하고 DataLoader에 입력합니다.

주석 부분을 보시면 DataLoader의 입력을 확인할 수 있습니다. shuffle 여부나 batch 크기 등을 정할 수 있습니다.

이번 코드에서는 아무런 설정없이 dataset만 입력하여 사용하였습니다.

dataDir = '/hostDir/dataSet'
annoVal = os.path.join(data_dir, 'annotations', 'person_keypoints_val2014.json')

# DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
#            batch_sampler=None, num_workers=0, collate_fn=None,
#            pin_memory=False, drop_last=False, timeout=0,
#            worker_init_fn=None, *, prefetch_factor=2,
#            persistent_workers=False)

myCocoDataset = MyCocoDataset(dataDir, annoVal)
myLoader = torch.utils.data.DataLoader(myCocoDataset)

출력:

loading annotations into memory...
Done (t=3.21s)
creating index...
index created!

 

 

 

DataLoader로 마지막으로 영상과 bbox들이 잘 출력되는지 확인하겠습니다.

tensor 형의 영상을 PIL로 변환하기 위해 to_pil_image를 사용하였습니다.

아래 결과를 보시면 원하는대로 영상과 bbox들이 출력되는 것을 확인할 수 있습니다.

from torchvision.transforms.functional import to_pil_image

for i, (img, bboxes) in enumerate(myLoader):
    print(f'index: {i}')
    plt.figure()
    plt.imshow(to_pil_image(img[0]))
    print(f'bboxes: \n{bboxes}\n')
    if i > 0:
        break

출력:

index: 0
bboxes: 
tensor([[[250.9000,  67.1000, 155.4200, 277.2300],
         [225.5600, 251.3600,  29.1700,  96.5400],
         [203.3800, 246.9900,  30.5000, 101.2600],
         [149.1400, 252.3500,  21.3700,  95.7500],
         [170.2800, 255.4200,  33.3700,  92.0500],
         [420.2600, 254.1900,  28.8500,  88.6700],
         [451.8300, 251.0700,  28.7800,  91.9900],
         [390.4400, 257.6100,  25.3900,  87.5900],
         [294.3500, 251.1500,  31.6400,  93.8400],
         [320.5400, 251.6900,  20.3500,  82.4600],
         [477.2800, 262.5500,  14.3700,  79.0000],
         [357.5500, 305.3700,  31.9400,  44.0100],
         [253.5200, 245.3700,  18.7800, 100.1600],
         [378.0000, 247.0000, 106.0000, 105.0000]]], dtype=torch.float64)

index: 1
bboxes: 
tensor([[[ 94.3000, 193.6700,  59.7400, 143.7500]]], dtype=torch.float64)

Index 0 영상

Index 1 영상

 

 

감사합니다.