[Pytorch] ImageFolder로 쉽게 Image Dataset 만들기

이미지를 Input으로 하는 네트워크를 구축하고자 할 때 유용하게 사용할 수 있는 torchvision의 Image Folder 클래스를 소개한다.

 

 

ImageFolder를 이용하면 이미지들이 각 class에 해당하는 폴더에 각각 나누어져 있을 때 이를 dataset 형태로 쉽게 불러올 수 있다.

예시를 위해 kaggle에서 제공하는 Medical MNIST 데이터를 살펴보겠다.

 

 

Medical MNIST는 다음과 같은 2차원 medical 이미지들이 총 6개의 class로 나뉘어져 있는 형태이다.

 

 

각 class의 이미지들은 위와 같이 class 이름으로 된 폴더에 나뉘어져 있다. 이를 ImageFolder를 이용해 데이터셋 형태로 만들어 보자.

우선 ImageFolder의 파라미터를 살펴보면 다음과 같다.

 

가장 중요한 파라미터는 이미지 파일의 경로를 알려주는 root이다. class별 폴더가 저장되어 있는 경로를 넣어 주면 된다.

이미지 처리를 위해서는 transform, target_transform (label도 같이 transform되어야 할 경우에 이용)을 사용한다.

이외에 이미지 로드 방식을 결정하는 loader, valid한 파일인지를 검사하는 is_valid_file 파라미터도 존재한다.

 

이미지를 불러오고 transform을 적용한 뒤 dataset으로 만드는 예시는 다음과 같다.

 

from torchvision import transforms
from torchvision.datasets import ImageFolder

# transform 정의 : Resize 후 Tensor 형태로
transform = transforms.Compose([
                                transforms.Resize((150, 150)),
                                transforms.ToTensor(),
])

data = ImageFolder(root = "./datasets/medMNIST", # 폴더의 root 경로 넣어주기
                   transform = transforms.ToTensor())

 

생성한 데이터셋은 바로 DataLoader를 통해 불러들일 수 있다!